fix: main.py — Экзамен: RAG-агент с ChromaDB и веб-поиском

main.py

@@ -1,136 +1,19 @@
-import os
 import asyncio
-from pathlib import Path
-
+import os
 from dotenv import load_dotenv
-from langchain_openai import ChatOpenAI, OpenAIEmbeddings
-from langchain_chroma import Chroma
-from langchain_core.documents import Document
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain.tools import tool
-from langchain_community.tools.tavily_search import TavilySearchResults
-
-from deepagents import create_deep_agent
-from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend
-from langchain_core.messages import HumanMessage
+from agent import run_agent
 
 load_dotenv()
 
-# ---------- LLM ----------
-llm = ChatOpenAI(
-    model="openai/gpt-oss-20b:free",
-    base_url="https://openrouter.ai/api/v1",
-    api_key=os.getenv("OPENAI_API_KEY"),
-    temperature=0.0,
-)
-
-# ---------- Vector Store ----------
-embeddings = OpenAIEmbeddings(
-    model="text-embedding-3-small",
-    base_url="https://openrouter.ai/api/v1",
-    api_key=os.getenv("OPENAI_API_KEY"),
-)
-
-def create_vectorstore(persist_directory: str = "./chroma_db") -> Chroma:
-    return Chroma(
-        collection_name="knowledge",
-        embedding_function=embeddings,
-        persist_directory=persist_directory,
-    )
-
-vectorstore = create_vectorstore()
-
-def load_documents(directory: str, vectorstore: Chroma) -> None:
-    path = Path(directory)
-    if not path.is_dir():
-        return
-    splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
-    docs = []
-    for file_path in path.rglob("*"):
-        if file_path.suffix.lower() not in {".txt", ".md"}:
-            continue
-        text = file_path.read_text(encoding="utf-8")
-        chunks = splitter.split_text(text)
-        for i, chunk in enumerate(chunks):
-            metadata = {"source": str(file_path), "chunk_index": i}
-            docs.append(Document(page_content=chunk, metadata=metadata))
-    if docs:
-        vectorstore.add_documents(docs)
-
-# Load initial documents (run once or each start)
-load_documents("./documents", vectorstore)
-
-# ---------- Tools ----------
-@tool
-def search_local_kb(query: str, top_k: int = 3) -> str:
-    """
-    Perform a semantic search in the local Chroma knowledge base.
-    Returns the concatenated contents of the most relevant documents.
-    """
-    docs = vectorstore.similarity_search(query, k=top_k)
-    if not docs:
-        return "No relevant information found in the local knowledge base."
-    return "\n---\n".join(doc.page_content for doc in docs)
-
-tavily_tool = TavilySearchResults(max_results=5)
-
-@tool
-def web_search(query: str) -> str:
-    """
-    Search the web using Tavily and return a short summary of the top results.
-    """
-    results = tavily_tool.run(query)
-    if not results:
-        return "No web results found."
-    # results is a list of dicts; extract title and url
-    lines = []
-    for r in results:
-        title = r.get("title", "No title")
-        url = r.get("url", "")
-        lines.append(f"{title}: {url}")
-    return "\n".join(lines)
-
-# ---------- Backend ----------
-backend = CompositeBackend(
-    [
-        LocalShellBackend(workspace_dir="./workspace"),
-        FilesystemBackend(),
-    ]
-)
-
-# ---------- Agent ----------
-system_prompt = (
-    "You are an AI assistant that can answer questions using two sources: "
-    "a local knowledge base (accessed via the tool `search_local_kb`) and the web (accessed via the tool `web_search`). "
-    "Decide which tool to use based on the user query. "
-    "When you use a tool, include the source name in your final answer: "
-    "`Source: chromadb` for local knowledge, `Source: tavily` for web results. "
-    "If both sources are needed, combine them and list both sources."
-)
-
-agent = create_deep_agent(
-    model=llm,
-    tools=[search_local_kb, web_search],
-    backend=backend,
-    system_prompt=system_prompt,
-)
-
-# ---------- CLI ----------
-async def chat_loop() -> None:
-    print("AI Assistant (type 'exit' to quit)")
-    thread_id = "session-1"
+async def main():
+    print("Введите запрос (или exit для выхода):")
     while True:
-        user_input = input("\nYou: ").strip()
-        if user_input.lower() in {"exit", "quit"}:
-            print("Goodbye!")
+        query = input("Запрос: ")
+        if query.lower() in ("exit", "quit"):
+            print("Выход.")
             break
-
-        result = await agent.ainvoke(
-            {"messages": [HumanMessage(content=user_input)]},
-            {"configurable": {"thread_id": thread_id}},
-        )
-        answer = result["messages"][-1].content
-        print(f"\nAssistant: {answer}")
+        response = await run_agent(query)
+        print(response)
 
 if __name__ == "__main__":
-    asyncio.run(chat_loop())
+    asyncio.run(main())