fix(needs_fixes): 1 исправлений, 0 отстояно — main.py

main.py

@@ -1,13 +1,11 @@
 import os
 import asyncio
-from pathlib import Path
 from dotenv import load_dotenv
-
-from langchain_openai import ChatOpenAI, OpenAIEmbeddings
+from langchain_ollama import ChatOllama
+from langchain_ollama import OllamaEmbeddings
 from langchain_chroma import Chroma
 from langchain_core.documents import Document
 from langchain_text_splitters import RecursiveCharacterTextSplitter
-from langchain_tavily import TavilySearchResults
 from langchain.tools import tool
 from deepagents import create_deep_agent
 from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend
@@ -15,97 +13,92 @@ from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeB
 # Load environment variables
 load_dotenv()
 
-# ---------- LLM and Embeddings ----------
-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,
-)
-
-embeddings = OpenAIEmbeddings(
-    model="text-embedding-3-small",
-    base_url="https://openrouter.ai/api/v1",
-    api_key=os.getenv("OPENAI_API_KEY"),
-)
-
 # ---------- Vector Store ----------
-CHROMA_DIR = Path("./chroma_db")
-vector_store = Chroma(
-    collection_name="knowledge",
-    embedding_function=embeddings,
-    persist_directory=str(CHROMA_DIR),
-)
+
+def create_vectorstore(persist_directory="./chroma_db"):
+    """Create a Chroma vector store with Ollama embeddings."""
+    embeddings = OllamaEmbeddings(model="nomic-embed-text")
+    return Chroma(persist_directory=persist_directory, embedding_function=embeddings)
+
+
+def load_documents(directory, vectorstore):
+    """Load .txt and .md files from *directory*, chunk them, and add to *vectorstore*.
+    The function preserves the file name in metadata for later reference.
+    """
+    splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+    for root, _, files in os.walk(directory):
+        for fname in files:
+            if fname.lower().endswith(('.txt', '.md')):
+                path = os.path.join(root, fname)
+                with open(path, 'r', encoding='utf-8') as f:
+                    text = f.read()
+                docs = splitter.split_text(text)
+                documents = [Document(page_content=chunk, metadata={"source": fname}) for chunk in docs]
+                vectorstore.add_documents(documents)
 
 # ---------- Tools ----------
+
+vectorstore = create_vectorstore()
+
 @tool
 def search_local_kb(query: str, top_k: int = 3) -> str:
-    """Semantic search in the local knowledge base."""
-    docs = vector_store.similarity_search(query, k=top_k)
+    """Semantic search in the local knowledge base (ChromaDB)."""
+    retriever = vectorstore.as_retriever(search_kwargs={"k": top_k})
+    docs = retriever.get_relevant_documents(query)
     if not docs:
-        return "No relevant documents found in local KB."
-    return "\n---\n".join(f"{i+1}. {doc.page_content[:200]}..." for i, doc in enumerate(docs))
+        return "No relevant local knowledge found."
+    return "\n---\n".join([f"{doc.metadata.get('source', 'unknown')}\n{doc.page_content}" for doc 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)
+    from tavily import TavilyClient
+    client = TavilyClient(api_key=os.getenv("TAVILY_API_KEY"))
+    results = client.search(query, max_results=3)
     if not results:
         return "No web results found."
-    return "\n---\n".join(f"{i+1}. {res['title']}\n{res['content'][:200]}..." for i, res in enumerate(results))
+    return "\n---\n".join([f"{r['title']}\n{r['content']}" for r in results])
+
+# ---------- Agent ----------
+
+llm = ChatOllama(model="llama3", temperature=0.0)
 
-# ---------- Backend ----------
 backend = CompositeBackend([
     LocalShellBackend(workspace_dir="./workspace"),
     FilesystemBackend(),
 ])
 
-# ---------- Agent ----------
+system_prompt = (
+    "You are an AI assistant with access to two tools: "
+    "search_local_kb for local knowledge and web_search for up-to-date information. "
+    "When answering a user query, first decide which tool is appropriate. "
+    "If the answer can be derived from the local documents, use search_local_kb; "
+    "otherwise use web_search. "
+    "Always indicate the source of the information in your response: "
+    "[Local KB] or [Web Search]."
+)
+
 agent = create_deep_agent(
     model=llm,
     tools=[search_local_kb, web_search],
     backend=backend,
-    system_prompt=(
-        "You are a helpful assistant. For any user query, decide whether to use the local knowledge base or perform a web search. "
-        "If the query is about recent events, news, or requires up‑to‑date information, use the web_search tool. "
-        "Otherwise, use search_local_kb. "
-        "Always return the source used in the response (either 'chromadb' or 'tavily')."
-    ),
+    system_prompt=system_prompt,
 )
 
-# ---------- Document Loader ----------
-def load_documents(directory: str, vectorstore: Chroma):
-    """Load .txt and .md files from a directory, chunk them, and add to the vector store."""
-    splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
-    docs = []
-    for file_path in Path(directory).glob("**/*"):
-        if file_path.suffix.lower() in {".txt", ".md"}:
-            text = file_path.read_text(encoding="utf-8")
-            chunks = splitter.split_text(text)
-            docs.extend([Document(page_content=chunk, metadata={"source": str(file_path)}) for chunk in chunks])
-    if docs:
-        vectorstore.add_documents(docs)
-        vectorstore.persist()
-
-# ---------- CLI ----------
 async def main():
-    # Ensure vector store is loaded
-    if not CHROMA_DIR.exists() or not any(CHROMA_DIR.iterdir()):
-        print("Loading documents into ChromaDB…")
-        load_documents("./documents", vector_store)
-    print("RAG agent ready. Type 'exit' to quit.")
+    print("RAG Agent ready. Type 'exit' to quit.")
     while True:
         user_input = input("\nЗапрос: ")
-        if user_input.lower() in {"exit", "quit"}:
+        if user_input.strip().lower() == "exit":
+            print("Goodbye!")
             break
         result = await agent.ainvoke(
             {"messages": [{"role": "user", "content": user_input}]},
             {"configurable": {"thread_id": "session-1"}},
         )
-        # Extract last message content
-        content = result["messages"][-1].content
-        print(f"\nОтвет:\n{content}")
+        # The last message is the assistant's reply
+        reply = result["messages"][-1].content
+        print(f"\n{reply}")
 
 if __name__ == "__main__":
     asyncio.run(main())