add main.py

main.py

@@ -1,18 +1,19 @@
 import os
 import asyncio
 from pathlib import Path
+from typing import List
+
 from langchain_openai import ChatOpenAI
 from langchain_core.messages import HumanMessage
 from langchain.tools import tool
+from langchain_ollama import OllamaEmbeddings
+from langchain_chroma import Chroma
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+from langchain_tavily import TavilySearchResults
 from deepagents import create_deep_agent
 from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend
-from langchain_chroma import Chroma
-from langchain_ollama import OllamaEmbeddings
-from langchain_text_splitters import RecursiveCharacterTextSplitter
-from langchain_ollama import OllamaEmbeddings
-from langchain_tavily import TavilySearchResults
 
-# ---------- LLM ----------
+# --------------------------- LLM ---------------------------
 llm = ChatOpenAI(
     model="openai/gpt-oss-20b:free",
     base_url="https://openrouter.ai/api/v1",
@@ -20,60 +21,55 @@ llm = ChatOpenAI(
     temperature=0.0,
 )
 
-# ---------- Backend ----------
+# --------------------------- Backend ---------------------------
 backend = CompositeBackend([
     LocalShellBackend(workspace_dir="./workspace"),
     FilesystemBackend(),
 ])
 
-# ---------- Vectorstore utilities ----------
+# --------------------------- Vectorstore ---------------------------
 PERSIST_DIR = Path("./chroma_db")
+PERSIST_DIR.mkdir(parents=True, exist_ok=True)
 
-def create_vectorstore(persist_directory: str = "./chroma_db"):
-    embeddings = OllamaEmbeddings(model="nomic-embed-text")
-    return Chroma(persist_directory=persist_directory, embedding_function=embeddings)
+embeddings = OllamaEmbeddings(model="nomic-embed-text")
+vectorstore = Chroma(persist_directory=str(PERSIST_DIR), embedding_function=embeddings)
 
-
-def load_documents(directory: str, vectorstore):
-    splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
-    docs = []
-    for file in Path(directory).glob("**/*"):
+# Load documents from ./documents if not already loaded
+DOCS_DIR = Path("./documents")
+if DOCS_DIR.exists():
+    for file in DOCS_DIR.glob("**/*.*"):
         if file.suffix.lower() in {".txt", ".md"}:
             text = file.read_text(encoding="utf-8")
-            docs.extend(splitter.split_text(text))
-    # Convert to Document objects
-    from langchain_core.documents import Document
-    documents = [Document(page_content=chunk) for chunk in docs]
-    vectorstore.add_documents(documents)
+            splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+            docs = splitter.split_text(text)
+            vectorstore.add_texts(docs, metadatas=[{"source": str(file)} for _ in docs])
     vectorstore.persist()
 
-# ---------- Tools ----------
+# --------------------------- Tools ---------------------------
 @tool
 def search_local_kb(query: str, top_k: int = 3) -> str:
     """Semantic search in the local ChromaDB knowledge base."""
-    vectorstore = create_vectorstore(PERSIST_DIR)
     retriever = vectorstore.as_retriever(search_kwargs={"k": top_k})
     docs = retriever.invoke(query)
     if not docs:
-        return "No relevant local documents found."
-    return "\n---\n".join(doc.page_content for doc in docs) + "\n[Source: chromadb]"
+        return "No local knowledge found."
+    return "\n---\n".join([f"{d.page_content[:500]}..." for d in docs])
 
 @tool
 def web_search(query: str) -> str:
     """Web search using Tavily."""
     tavily = TavilySearchResults(api_key=os.getenv("TAVILY_API_KEY"))
-    results = tavily.run(query)
+    results = tavily.invoke(query)
     if not results:
         return "No web results found."
-    snippets = [f"{r['title']}\n{r['content']}" for r in results]
-    return "\n---\n".join(snippets) + "\n[Source: tavily]"
+    return "\n---\n".join([f"{r['title']}: {r['content'][:500]}..." for r in results])
 
-# ---------- Agent ----------
+# --------------------------- Agent ---------------------------
 SYSTEM_PROMPT = (
     "You are an AI assistant that can answer questions using either a local knowledge base or the web. "
-    "If the question refers to documents in the local folder, use the `search_local_kb` tool. "
-    "If the question is about recent events or requires up‑to‑date information, use the `web_search` tool. "
-    "Always include the source tag (`[Source: chromadb]` or `[Source: tavily]`) in your answer."
+    "If the answer can be found in the local documents, use the `search_local_kb` tool and prefix the response with `[Local KB]`. "
+    "If the answer requires up‑to‑date information, use the `web_search` tool and prefix the response with `[Web Search]`. "
+    "Always indicate the source in the response."
 )
 
 agent = create_deep_agent(
@@ -83,12 +79,12 @@ agent = create_deep_agent(
     system_prompt=SYSTEM_PROMPT,
 )
 
-# ---------- CLI ----------
-async def chat_loop():
+# --------------------------- CLI ---------------------------
+async def main():
     print("RAG Agent ready. Type 'exit' to quit.")
     while True:
         user_input = input("\nЗапрос: ")
-        if user_input.strip().lower() == "exit":
+        if user_input.lower() in {"exit", "quit"}:
             print("Goodbye!")
             break
         result = await agent.ainvoke(
@@ -97,14 +93,7 @@ async def chat_loop():
         )
         # The last message is the assistant's reply
         reply = result["messages"][-1].content
-        print(reply)
+        print(f"\n{reply}")
 
-# ---------- Initialization ----------
 if __name__ == "__main__":
-    # Ensure vectorstore exists and load documents if empty
-    vectorstore = create_vectorstore(PERSIST_DIR)
-    if not vectorstore.get_all_documents():
-        print("Loading documents into ChromaDB...")
-        load_documents("./documents", vectorstore)
-        print("Documents loaded.")
-    asyncio.run(chat_loop())
+    asyncio.run(main())