diff --git a/main.py b/main.py index 31f3481..6934782 100644 --- a/main.py +++ b/main.py @@ -1,21 +1,20 @@ -import os, asyncio +import os +import asyncio from dotenv import load_dotenv -from langchain_openai import ChatOpenAI -from langchain_core.messages import HumanMessage +from langchain_openai import ChatOpenAI, OpenAIEmbeddings +from langchain_chroma import Chroma +from langchain_text_splitters import RecursiveCharacterTextSplitter +from langchain_core.documents import Document +from langchain_tavily import TavilySearchResults from langchain.tools import tool 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 +from langchain_core.messages import HumanMessage -# Load env vars +# Load environment variables load_dotenv() -# ---------- LLM ---------- +# ---------- LLM and Embeddings ---------- llm = ChatOpenAI( model="openai/gpt-oss-20b:free", base_url="https://openrouter.ai/api/v1", @@ -23,43 +22,60 @@ 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), +embeddings = OpenAIEmbeddings( + model="text-embedding-3-small", + base_url="https://openrouter.ai/api/v1", + api_key=os.getenv("OPENAI_API_KEY"), ) -# 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() +# ---------- Vector Store ---------- +CHROMA_DIR = "./chroma_db" +vector_store = Chroma(collection_name="knowledge", embedding_function=embeddings, persist_directory=CHROMA_DIR) + +# ---------- Document Loader ---------- + +def load_documents(directory: str): + """Read .txt/.md files, split into chunks, and add to Chroma collection.""" + splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200) + docs = [] + for root, _, files in os.walk(directory): + for file in files: + if file.lower().endswith(('.txt', '.md')): + path = os.path.join(root, file) + with open(path, 'r', encoding='utf-8') as f: + text = f.read() + chunks = splitter.split_text(text) + docs.extend([Document(page_content=c, metadata={"source": path}) for c in chunks]) + if docs: + vector_store.add_documents(docs) + vector_store.persist() + +# Load documents once at startup +if not os.path.exists(CHROMA_DIR) or not os.listdir(CHROMA_DIR): + load_documents("./documents") + +# ---------- Tavily Search ---------- +TAVILY_API_KEY = os.getenv("TAVILY_API_KEY") +if not TAVILY_API_KEY: + raise ValueError("TAVILY_API_KEY not set in .env") # ---------- 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) + """Semantic search in the local Chroma knowledge base.""" + docs = vector_store.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: + """Web search using Tavily.""" + results = TavilySearchResults(api_key=TAVILY_API_KEY, max_results=3) + search_results = results.run(query) + if not search_results: return "No web results found." - return "\n---\n".join([f"{r['title']}\n{r['content']}" for r in results]) + return "\n---\n".join([f"{r['title']}\n{r['content']}" for r in search_results]) # ---------- Backend ---------- backend = CompositeBackend([ @@ -72,13 +88,21 @@ 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 your answer.", + system_prompt=( + "You are a RAG agent with a local knowledge base and web search capability. " + "When a user asks a question, decide whether the answer can be found in the local " + "knowledge base or requires up‑to‑date information from the web. Use the tool " + "search_local_kb for local queries and web_search for web queries. " + "Always indicate the source of the information in your final answer: " + "(chromadb) or (tavily)." + ), ) +# ---------- CLI ---------- async def main(): - print("Welcome to the RAG agent. Type 'exit' to quit.") + print("RAG Agent ready. Type 'exit' to quit.") while True: - user_input = input("\nQuery: ") + user_input = input("\nЗапрос: ") if user_input.lower() in {"exit", "quit"}: print("Goodbye!") break @@ -86,9 +110,9 @@ async def main(): {"messages": [HumanMessage(content=user_input)]}, {"configurable": {"thread_id": "session-1"}}, ) - # The agent returns a list of messages; last is the assistant reply + # The agent returns a list of messages; the last is the assistant reply reply = result["messages"][-1].content - print("\nAnswer:\n", reply) + print(f"\n{reply}") if __name__ == "__main__": asyncio.run(main())