feat: solution for 'Агент с RAG-памятью'

src/main.py

@@ -1,110 +1,20 @@
+#!/usr/bin/env python3
 """
-Simple RAG agent using LangChain, Qdrant, and Ollama.
-
-This script demonstrates how to set up a retrieval-augmented generation (RAG) pipeline
-with a local Qdrant vector store and an Ollama LLM. It can be run directly:
-
-    python -m src.main
-
-The script will prompt the user for a question and return an answer based on the
-documents stored in Qdrant.
-
-Prerequisites:
-- Qdrant server running locally (default port 6333).
-- Ollama server running locally (default port 11434).
-- A Qdrant collection named "rag_collection" populated with embeddings.
+Main entry point for the knowledge‑base agent.
 """
 
-import os
-import sys
-from typing import Optional
-
-try:
-    from langchain_ollama import OllamaLLM
-    from langchain_qdrant import QdrantStore
-    from langchain.chains import RetrievalQA
-    from langchain.memory import ConversationBufferMemory
-except ImportError as e:
-    print("Required packages are missing. Please run 'pip install -r requirements.txt'.")
-    sys.exit(1)
-
-
-def get_llm() -> OllamaLLM:
-    """
-    Create an Ollama LLM instance.
-    """
-    # Ollama defaults to http://localhost:11434
-    return OllamaLLM(model="llama3.1")
-
-
-def get_vector_store() -> QdrantStore:
-    """
-    Connect to the local Qdrant instance and load the collection.
-    """
-    # Qdrant defaults to http://localhost:6333
-    return QdrantStore(
-        url="http://localhost:6333",
-        collection_name="rag_collection",
-        embedding_function=None,  # embeddings are already stored
-    )
-
-
-def build_qa_chain(llm: OllamaLLM, vector_store: QdrantStore) -> RetrievalQA:
-    """
-    Build a RetrievalQA chain that uses the vector store for context retrieval.
-    """
-    memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
-
-    return RetrievalQA.from_chain_type(
-        llm=llm,
-        chain_type="stuff",
-        retriever=vector_store.as_retriever(search_kwargs={"k": 4}),
-        memory=memory,
-        return_source_documents=True,
-    )
+from .knowledge_base import KnowledgeBase
+from .tools.knowledge_base_tool import KnowledgeBaseTool
+from .cli import run_cli
 
 
 def main() -> None:
     """
-    Main entry point: prompt user for a question and print the answer.
+    Create the knowledge base, wrap it in a tool, and start the CLI.
     """
-    print("Initializing RAG agent...")
-    try:
-        llm = get_llm()
-        vector_store = get_vector_store()
-        qa_chain = build_qa_chain(llm, vector_store)
-    except Exception as exc:
-        print(f"Failed to initialize components: {exc}")
-        sys.exit(1)
-
-    print("RAG agent ready. Type your question (or 'exit' to quit).")
-    while True:
-        try:
-            user_input = input("\n> ").strip()
-        except (EOFError, KeyboardInterrupt):
-            print("\nExiting.")
-            break
-
-        if user_input.lower() in {"exit", "quit"}:
-            print("Goodbye!")
-            break
-
-        if not user_input:
-            print("Please enter a non-empty question.")
-            continue
-
-        try:
-            result = qa_chain({"question": user_input})
-            answer = result.get("answer", "No answer returned.")
-            sources = result.get("source_documents", [])
-            print("\nAnswer:")
-            print(answer)
-            if sources:
-                print("\nSources:")
-                for doc in sources:
-                    print(f"- {doc.metadata.get('source', 'unknown')}")
-        except Exception as exc:
-            print(f"Error during query: {exc}")
+    kb = KnowledgeBase()
+    kb_tool = KnowledgeBaseTool(kb)
+    run_cli(kb_tool)
 
 
 if __name__ == "__main__":