From f64bf247752dcb6eb83e9af8fa8443bd6b99fe74 Mon Sep 17 00:00:00 2001 From: Danil Parunin 5f1b81b8-4f5d-11e8-9c2d-fa7ae01bbebc Date: Tue, 16 Jun 2026 15:58:30 +0000 Subject: [PATCH] Add main.py --- main.py | 206 ++++++++++++++++---------------------------------------- 1 file changed, 58 insertions(+), 148 deletions(-) diff --git a/main.py b/main.py index 47eeb16..b0341de 100644 --- a/main.py +++ b/main.py @@ -1,163 +1,73 @@ import os -import asyncio import json +import httpx from pathlib import Path -from typing import List - -from langchain_openai import ChatOpenAI, OpenAIEmbeddings +from langchain_ollama import ChatOllama from langchain_chroma import Chroma -from langchain_core.documents import Document -from langchain_core.messages import HumanMessage -from langchain.tools import tool -from deepagents import create_deep_agent -from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend +from langchain.embeddings import OllamaEmbeddings +from langchain.schema import Document +from langchain.prompts import ChatPromptTemplate +from langchain.chains import LLMChain -# --------------------------------------------------------------------------- -# Configuration -# --------------------------------------------------------------------------- -OPENAI_API_KEY = os.getenv("OPENAI_API_KEY") -if not OPENAI_API_KEY: - raise RuntimeError("OPENAI_API_KEY not set in environment") +# 1. Load FAQ into Chroma -# LLM via OpenRouter -llm = ChatOpenAI( - model="openai/gpt-oss-20b:free", - base_url="https://openrouter.ai/api/v1", - api_key=OPENAI_API_KEY, - temperature=0.0, -) +def load_faq_to_chroma(md_path: str, persist_dir: str = "./chroma_faq"): + from langchain.document_loaders import TextLoader + from langchain.text_splitter import RecursiveCharacterTextSplitter -# Embeddings via OpenRouter -embeddings = OpenAIEmbeddings( - model="text-embedding-3-small", - base_url="https://openrouter.ai/api/v1", - api_key=OPENAI_API_KEY, -) + loader = TextLoader(md_path) + docs = loader.load() + splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50) + texts = splitter.split_documents(docs) + embeddings = OllamaEmbeddings(model="nomic-embed-text") + chroma = Chroma.from_documents(texts, embeddings, persist_directory=persist_dir) + chroma.persist() + return chroma -# --------------------------------------------------------------------------- -# 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), -) +# 2. Search function -# --------------------------------------------------------------------------- -# Utility: load markdown files into Chroma -# --------------------------------------------------------------------------- -@tool -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." +def search_course_docs(query: str, k: int = 3): + chroma = Chroma(persist_directory="./chroma_faq", embedding_function=OllamaEmbeddings(model="nomic-embed-text")) + results = chroma.similarity_search(query, k=k) + return [doc.page_content for doc in results] -# --------------------------------------------------------------------------- -# 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}" +# 3. MCP-style tool -# --------------------------------------------------------------------------- -# MCP-style tool: fetch course metadata -# --------------------------------------------------------------------------- -# For simplicity we use a static JSON file in the repo. In production this -# would be an HTTP call to an MCP server. -METADATA_FILE = Path("course_meta.json") +def fetch_course_meta(query: str): + # For demo, use static JSON file + meta_path = Path("meta.json") + if not meta_path.exists(): + return {"error": "Meta not found"} + meta = json.loads(meta_path.read_text()) + # simple search by key + return meta.get(query, {}) -@tool -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) +# 4. Agent logic -# --------------------------------------------------------------------------- -# Backend setup -# --------------------------------------------------------------------------- -backend = CompositeBackend([ - LocalShellBackend(workspace_dir="./workspace"), - FilesystemBackend(), -]) - -# --------------------------------------------------------------------------- -# 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) +def answer_question(question: str): + # Simple heuristic: if question contains "schedule" or "метаданные" use meta + if any(word in question.lower() for word in ["schedule", "расписание", "метаданные"]): + source = "mcp_meta" + answer = fetch_course_meta(question) + else: + source = "chroma" + answer = search_course_docs(question, k=1)[0] + return {"answer": answer, "source": source} 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']})")