diff --git a/main.py b/main.py index 997f136..a13c071 100644 --- a/main.py +++ b/main.py @@ -1,19 +1,22 @@ import os import asyncio from pathlib import Path -from typing import List +from dotenv import load_dotenv 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_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 -# --------------------------- LLM --------------------------- +# Load env vars +load_dotenv() + +# ---------- LLM ---------- llm = ChatOpenAI( model="openai/gpt-oss-20b:free", base_url="https://openrouter.ai/api/v1", @@ -21,79 +24,65 @@ llm = ChatOpenAI( temperature=0.0, ) -# --------------------------- Backend --------------------------- +# ---------- Backend ---------- backend = CompositeBackend([ LocalShellBackend(workspace_dir="./workspace"), FilesystemBackend(), ]) -# --------------------------- Vectorstore --------------------------- +# ---------- Vector Store ---------- PERSIST_DIR = Path("./chroma_db") -PERSIST_DIR.mkdir(parents=True, exist_ok=True) +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 -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") - 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() +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 --------------------------- +# ---------- Tools ---------- @tool def search_local_kb(query: str, top_k: int = 3) -> str: - """Semantic search in the local ChromaDB knowledge base.""" + """Semantic search in local ChromaDB knowledge base.""" retriever = vectorstore.as_retriever(search_kwargs={"k": top_k}) docs = retriever.invoke(query) - if not docs: - return "No local knowledge found." - return "\n---\n".join([f"{d.page_content[:500]}..." for d in docs]) + 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 using Tavily.""" + """Web search via Tavily.""" tavily = TavilySearchResults(api_key=os.getenv("TAVILY_API_KEY")) results = tavily.invoke(query) - if not results: - return "No web results found." - return "\n---\n".join([f"{r['title']}: {r['content'][:500]}..." for r in results]) - -# --------------------------- Agent --------------------------- -SYSTEM_PROMPT = ( - "You are an AI assistant that can answer questions using either a local knowledge base or the web. " - "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." -) + 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=SYSTEM_PROMPT, + 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.", ) -# --------------------------- CLI --------------------------- +# ---------- CLI ---------- async def main(): print("RAG Agent ready. Type 'exit' to quit.") while True: user_input = input("\nЗапрос: ") if user_input.lower() in {"exit", "quit"}: - print("Goodbye!") break result = await agent.ainvoke( {"messages": [HumanMessage(content=user_input)]}, {"configurable": {"thread_id": "session-1"}}, ) - # The last message is the assistant's reply + # The agent returns a list of messages; last is the assistant reply reply = result["messages"][-1].content - print(f"\n{reply}") + print(f"\nОтвет: {reply}") if __name__ == "__main__": asyncio.run(main())