diff --git a/main.py b/main.py index ee6117e..6fa6559 100644 --- a/main.py +++ b/main.py @@ -1,105 +1,76 @@ -import os -import asyncio -from pathlib import Path -from dotenv import load_dotenv -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.tools import tool -from deepagents import create_deep_agent -from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend +"""RAG-агент с ChromaDB (Ollama) и веб-поиском (Tavily). CLI.""" +from __future__ import annotations + +from pathlib import Path + +from dotenv import load_dotenv +from langchain_ollama import ChatOllama +from langchain.agents import AgentExecutor, create_react_agent +from langchain import hub + +from vectorstore import create_vectorstore, load_documents +from tools import init_tools, search_local_kb, web_search -# 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, +# ── LLM: Ollama llama3 ──────────────────────────────────────────────────── +llm = ChatOllama(model="llama3", temperature=0) + +# ── Векторное хранилище ─────────────────────────────────────────────────── +CHROMA_DIR = "./chroma_db" +vectorstore = create_vectorstore(CHROMA_DIR) +init_tools(vectorstore) + +# Автозагрузка документов из documents/ (если есть и БД пустая) +docs_dir = Path("documents") +if docs_dir.exists() and not any(Path(CHROMA_DIR).rglob("*.sqlite3")): + count = load_documents(str(docs_dir), vectorstore) + print(f"[init] Загружено {count} чанков из {docs_dir}/") + +# ── ReAct-агент ─────────────────────────────────────────────────────────── +tools = [search_local_kb, web_search] +prompt = hub.pull("hwchase17/react") +agent = create_react_agent(llm, tools, prompt) +agent_executor = AgentExecutor( + agent=agent, + tools=tools, + verbose=True, + max_iterations=10, + handle_parsing_errors=True, ) -embeddings = OpenAIEmbeddings( - model="text-embedding-3-small", - base_url="https://openrouter.ai/api/v1", - api_key=os.getenv("OPENAI_API_KEY"), +ROUTING_HINT = ( + "Для вопросов по локальным документам/конспектам используй search_local_kb. " + "Для актуальных новостей и фактов из интернета — web_search. " + "В ответе всегда укажи источник: chromadb или tavily." ) -# ---------- Vector Store ---------- -CHROMA_DIR = Path("./chroma_db") -vector_store = Chroma( - collection_name="knowledge", - embedding_function=embeddings, - persist_directory=str(CHROMA_DIR), -) -# ---------- Document Loader ---------- +def ask(question: str) -> str: + """Задать вопрос агенту.""" + result = agent_executor.invoke({"input": f"{ROUTING_HINT} -def load_documents(directory: str, vectorstore: Chroma): - splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200) - docs = [] - for file_path in Path(directory).glob("**/*.*"): - if file_path.suffix.lower() not in {".txt", ".md"}: - continue - text = file_path.read_text(encoding="utf-8") - chunks = splitter.split_text(text) - docs.extend([Document(page_content=c, metadata={"source": str(file_path)}) for c in chunks]) - if docs: - vectorstore.add_documents(docs) - vectorstore.persist() +Вопрос: {question}"}) + return result.get("output", str(result)) -# Load initial documents if collection is empty -if not CHROMA_DIR.exists() or not list(CHROMA_DIR.iterdir()): - load_documents("documents", vector_store) - -# ---------- Tools ---------- -@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) - if not docs: - return "No relevant information found in local knowledge base." - 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 using Tavily.""" - from langchain_tavily import TavilySearchResults - tavily = TavilySearchResults() - results = tavily.run(query) - return "\n---\n".join([f"{r['title']}\n{r['content']}" for r in results]) - -# ---------- Backend ---------- -backend = CompositeBackend([ - LocalShellBackend(workspace_dir="./workspace"), - FilesystemBackend(), -]) - -# ---------- Agent ---------- -agent = create_deep_agent( - model=llm, - tools=[search_local_kb, web_search], - backend=backend, - system_prompt="You are a helpful assistant that can search both a local knowledge base and the web.\nWhen answering, always indicate the source: either 'chromadb' or 'tavily'.\nUse the appropriate tool based on the query context.", -) - -# ---------- CLI ---------- -async def main(): - print("RAG Agent with ChromaDB and Tavily. 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": [{"role": "user", "content": user_input}]}, - {"configurable": {"thread_id": "session-1"}}, - ) - # Extract last message content - content = result["messages"][-1]["content"] - print(content) if __name__ == "__main__": - asyncio.run(main()) + print("RAG-агент: Ollama llama3 + nomic-embed-text + ChromaDB + Tavily") + print("Введите 'exit' для выхода.") + print("-" * 60) + while True: + try: + user_input = input(" +Запрос: ").strip() + except (EOFError, KeyboardInterrupt): + print(" +Завершение.") + break + if user_input.lower() in ("exit", "quit", "выход"): + print("До свидания!") + break + if not user_input: + continue + answer = ask(user_input) + print(f" +{answer}")