diff --git a/main.py b/main.py index a13c071..31f3481 100644 --- a/main.py +++ b/main.py @@ -1,17 +1,16 @@ -import os -import asyncio -from pathlib import Path +import os, asyncio from dotenv import load_dotenv - from langchain_openai import ChatOpenAI from langchain_core.messages import HumanMessage from langchain.tools import tool -from langchain_text_splitters import RecursiveCharacterTextSplitter -from langchain_chroma import Chroma -from langchain_ollama import OllamaEmbeddings -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_text_splitters import RecursiveCharacterTextSplitter +from langchain_ollama import OllamaEmbeddings +from langchain_ollama import Ollama +from langchain_tavily import TavilySearchResults +from pathlib import Path # Load env vars load_dotenv() @@ -24,57 +23,64 @@ llm = ChatOpenAI( temperature=0.0, ) +# ---------- Vectorstore ---------- +persist_dir = Path("./chroma_db") +persist_dir.mkdir(parents=True, exist_ok=True) + +embeddings = OllamaEmbeddings(model="nomic-embed-text") +vectorstore = Chroma( + collection_name="knowledge", + embedding_function=embeddings, + persist_directory=str(persist_dir), +) + +# Load documents from ./documents +text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200) +for file_path in Path("./documents").glob("**/*.*"): + if file_path.suffix.lower() in {".txt", ".md"}: + content = file_path.read_text(encoding="utf-8") + docs = text_splitter.split_text(content) + vectorstore.add_documents([{"page_content": d, "metadata": {"source": str(file_path)}} for d in docs]) +vectorstore.persist() + +# ---------- Tools ---------- +@tool +def search_local_kb(query: str, top_k: int = 3) -> str: + """Semantic search in local ChromaDB knowledge base.""" + docs = vectorstore.similarity_search(query, k=top_k) + if not docs: + return "No relevant local knowledge found." + return "\n---\n".join([f"{d.metadata.get('source', 'unknown')}\n{d.page_content}" for d in docs]) + +@tool +def web_search(query: str) -> str: + """Web search via Tavily.""" + tavily = TavilySearchResults(api_key=os.getenv("TAVILY_API_KEY")) + results = tavily.run(query) + if not results: + return "No web results found." + return "\n---\n".join([f"{r['title']}\n{r['content']}" for r in results]) + # ---------- Backend ---------- backend = CompositeBackend([ LocalShellBackend(workspace_dir="./workspace"), FilesystemBackend(), ]) -# ---------- Vector Store ---------- -PERSIST_DIR = Path("./chroma_db") -PERSIST_DIR.mkdir(exist_ok=True) - -embeddings = OllamaEmbeddings(model="nomic-embed-text") -vectorstore = Chroma(persist_directory=str(PERSIST_DIR), embedding_function=embeddings) - -# Load documents from ./documents if not already loaded -if not vectorstore.get_collection().count(): - splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200) - docs = [] - for file in Path("./documents").glob("*.txt"): - text = file.read_text(encoding="utf-8") - docs.extend(splitter.split_text(text)) - vectorstore.add_texts(docs) - -# ---------- Tools ---------- -@tool -def search_local_kb(query: str, top_k: int = 3) -> str: - """Semantic search in local ChromaDB knowledge base.""" - retriever = vectorstore.as_retriever(search_kwargs={"k": top_k}) - docs = retriever.invoke(query) - return "\n".join(doc.page_content for doc in docs) if docs else "No local results found." - -@tool -def web_search(query: str) -> str: - """Web search via Tavily.""" - tavily = TavilySearchResults(api_key=os.getenv("TAVILY_API_KEY")) - results = tavily.invoke(query) - return "\n".join(f"{r['title']}: {r['url']}" for r in results) if results else "No web results found." - # ---------- Agent ---------- agent = create_deep_agent( model=llm, tools=[search_local_kb, web_search], backend=backend, - system_prompt="You are a helpful RAG agent. For questions about local documents use search_local_kb, for up‑to‑date facts use web_search. Always state the source (chromadb or tavily) in the answer.", + system_prompt="You are a helpful RAG agent. For questions about local documents use search_local_kb, for up‑to‑date facts use web_search. Always state the source (chromadb or tavily) in your answer.", ) -# ---------- CLI ---------- async def main(): - print("RAG Agent ready. Type 'exit' to quit.") + print("Welcome to the RAG agent. Type 'exit' to quit.") while True: - user_input = input("\nЗапрос: ") + user_input = input("\nQuery: ") if user_input.lower() in {"exit", "quit"}: + print("Goodbye!") break result = await agent.ainvoke( {"messages": [HumanMessage(content=user_input)]}, @@ -82,7 +88,7 @@ async def main(): ) # The agent returns a list of messages; last is the assistant reply reply = result["messages"][-1].content - print(f"\nОтвет: {reply}") + print("\nAnswer:\n", reply) if __name__ == "__main__": asyncio.run(main())