Add main.py

main.py

@@ -1,163 +1,73 @@
 import os
-import asyncio
 import json
+import httpx
 from pathlib import Path
-from typing import List
+from langchain_ollama import ChatOllama
-
-from langchain_openai import ChatOpenAI, OpenAIEmbeddings
 from langchain_chroma import Chroma
-from langchain_core.documents import Document
+from langchain.embeddings import OllamaEmbeddings
-from langchain_core.messages import HumanMessage
+from langchain.schema import Document
-from langchain.tools import tool
+from langchain.prompts import ChatPromptTemplate
-from deepagents import create_deep_agent
+from langchain.chains import LLMChain
-from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend
 
-# ---------------------------------------------------------------------------
+# 1. Load FAQ into Chroma
-# Configuration
-# ---------------------------------------------------------------------------
-OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
-if not OPENAI_API_KEY:
-    raise RuntimeError("OPENAI_API_KEY not set in environment")
 
-# LLM via OpenRouter
+def load_faq_to_chroma(md_path: str, persist_dir: str = "./chroma_faq"):
-llm = ChatOpenAI(
+    from langchain.document_loaders import TextLoader
-    model="openai/gpt-oss-20b:free",
+    from langchain.text_splitter import RecursiveCharacterTextSplitter
-    base_url="https://openrouter.ai/api/v1",
-    api_key=OPENAI_API_KEY,
-    temperature=0.0,
-)
 
-# Embeddings via OpenRouter
+    loader = TextLoader(md_path)
-embeddings = OpenAIEmbeddings(
+    docs = loader.load()
-    model="text-embedding-3-small",
+    splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
-    base_url="https://openrouter.ai/api/v1",
+    texts = splitter.split_documents(docs)
-    api_key=OPENAI_API_KEY,
+    embeddings = OllamaEmbeddings(model="nomic-embed-text")
-)
+    chroma = Chroma.from_documents(texts, embeddings, persist_directory=persist_dir)
+    chroma.persist()
+    return chroma
 
-# ---------------------------------------------------------------------------
+# 2. Search function
-# Chroma persistence
-# ---------------------------------------------------------------------------
-CHROMA_DIR = Path("./chroma_faq")
-CHROMA_DIR.mkdir(parents=True, exist_ok=True)
-vector_store = Chroma(
-    collection_name="faq_collection",
-    embedding_function=embeddings,
-    persist_directory=str(CHROMA_DIR),
-)
 
-# ---------------------------------------------------------------------------
+def search_course_docs(query: str, k: int = 3):
-# Utility: load markdown files into Chroma
+    chroma = Chroma(persist_directory="./chroma_faq", embedding_function=OllamaEmbeddings(model="nomic-embed-text"))
-# ---------------------------------------------------------------------------
+    results = chroma.similarity_search(query, k=k)
-@tool
+    return [doc.page_content for doc in results]
-def load_faq_to_chroma() -> str:
-    """Load all .md files from data/ into the Chroma vector store.
-    This tool is idempotent – it will overwrite existing collection.
-    """
-    data_dir = Path("data")
-    if not data_dir.exists():
-        return "Data directory not found."
-    docs: List[Document] = []
-    for md_file in data_dir.glob("*.md"):
-        text = md_file.read_text(encoding="utf-8")
-        docs.append(Document(page_content=text, metadata={"source": md_file.name}))
-    if docs:
-        vector_store.delete_collection()
-        vector_store.add_documents(docs)
-        vector_store.persist()
-        return f"Loaded {len(docs)} documents into Chroma."
-    return "No markdown files found."
 
-# ---------------------------------------------------------------------------
+# 3. MCP-style tool
-# Tool: search knowledge base
-# ---------------------------------------------------------------------------
-@tool
-def search_course_docs(query: str, k: int = 3) -> str:
-    """Search the FAQ collection for relevant passages.
-    Returns a string with the top k passages and a source tag.
-    """
-    docs = vector_store.similarity_search(query, k=k)
-    if not docs:
-        return "No relevant information found in the FAQ."
-    results = "\n\n".join(f"{i+1}. {doc.page_content}" for i, doc in enumerate(docs))
-    return f"source: chroma\n{results}"
 
-# ---------------------------------------------------------------------------
+def fetch_course_meta(query: str):
-# MCP-style tool: fetch course metadata
+    # For demo, use static JSON file
-# ---------------------------------------------------------------------------
+    meta_path = Path("meta.json")
-# For simplicity we use a static JSON file in the repo. In production this
+    if not meta_path.exists():
-# would be an HTTP call to an MCP server.
+        return {"error": "Meta not found"}
-METADATA_FILE = Path("course_meta.json")
+    meta = json.loads(meta_path.read_text())
+    # simple search by key
+    return meta.get(query, {})
 
-@tool
+# 4. Agent logic
-def fetch_course_meta(query: str) -> str:
-    """Return metadata that matches the query.
-    The function performs a simple keyword search in the static JSON.
-    """
-    if not METADATA_FILE.exists():
-        return "Metadata file not found."
-    data = json.loads(METADATA_FILE.read_text(encoding="utf-8"))
-    matches = [item for item in data if query.lower() in item.get("title", "").lower()]
-    if not matches:
-        return "No metadata matches the query."
-    return f"source: mcp_meta\n" + json.dumps(matches, indent=2)
 
-# ---------------------------------------------------------------------------
+def answer_question(question: str):
-# Backend setup
+    # Simple heuristic: if question contains "schedule" or "метаданные" use meta
-# ---------------------------------------------------------------------------
+    if any(word in question.lower() for word in ["schedule", "расписание", "метаданные"]):
-backend = CompositeBackend([
+        source = "mcp_meta"
-    LocalShellBackend(workspace_dir="./workspace"),
+        answer = fetch_course_meta(question)
-    FilesystemBackend(),
+    else:
-])
+        source = "chroma"
-
+        answer = search_course_docs(question, k=1)[0]
-# ---------------------------------------------------------------------------
+    return {"answer": answer, "source": source}
-# Agent definition
-# ---------------------------------------------------------------------------
-agent = create_deep_agent(
-    model=llm,
-    tools=[load_faq_to_chroma, search_course_docs, fetch_course_meta],
-    backend=backend,
-    system_prompt=(
-        "You are a helpful FAQ bot for a course. "
-        "Use the search_course_docs tool for questions about lecture materials. "
-        "Use fetch_course_meta for questions about schedule or metadata. "
-        "Do not call both tools unless necessary. "
-        "Always prefix your answer with the source tag (chroma or mcp_meta)."
-    ),
-)
-
-# ---------------------------------------------------------------------------
-# CLI helpers
-# ---------------------------------------------------------------------------
-PRESET_QUESTIONS = [
-    "What is the deadline for the final project?",  # should hit metadata
-    "Explain the concept of polymorphism in OOP.",  # should hit FAQ
-    "How many lectures are there in the first module?",  # metadata
-]
-
-async def run_agent(question: str, thread_id: str = "session-1"):
-    result = await agent.ainvoke(
-        {"messages": [HumanMessage(content=question)]},
-        {"configurable": {"thread_id": thread_id}},
-    )
-    return result["messages"][-1].content
-
-async def main():
-    # Ensure FAQ is loaded once
-    await agent.ainvoke(
-        {"messages": [HumanMessage(content="load_faq_to_chroma()")]},
-        {"configurable": {"thread_id": "init"}},
-    )
-    print("\n--- Preset questions ---")
-    for q in PRESET_QUESTIONS:
-        ans = await run_agent(q)
-        print(f"Q: {q}\nA: {ans}\n")
-    print("\n--- Interactive mode (type 'exit' to quit) ---")
-    while True:
-        user_input = input("> ")
-        if user_input.lower() in {"exit", "quit"}:
-            break
-        ans = await run_agent(user_input)
-        print(ans)
 
 if __name__ == "__main__":
-    asyncio.run(main())
+    # CLI with preset questions
+    preset = [
+        "Что такое ChromaDB?",
+        "Как подключить Ollama embeddings?",
+        "Когда будет расписание следующего занятия?"
+    ]
+    for q in preset:
+        res = answer_question(q)
+        print(f"Q: {q}\nA: {res['answer']}\nSource: {res['source']}\n")
+    # interactive
+    while True:
+        q = input("Ask a question (or 'exit'): ")
+        if q.lower() == "exit":
+            break
+        res = answer_question(q)
+        print(f"A: {res['answer']} (source: {res['source']})")