feat: solution for 6a02e23da6fe2e4ac16acf65

solutions/6a02e23da6fe2e4ac16acf65/solution.py

@@ -1,95 +1,126 @@
-from langchain_ollama import OllamaEmbeddings
-from langchain_qdrant import QdrantVectorStore
-from qdrant_client import QdrantClient
-from qdrant_client.http.models import Distance, VectorParams
-from langchain_core.documents import Document
-from langchain_text_splitters import RecursiveCharacterTextSplitter
-from langchain.tools import tool
-from langchain.agents import create_agent
-from langchain_openai import ChatOpenAI
-from pydantic import SecretStr
+from pathlib import Path
 
-# ---------- LLM and embeddings ----------
-llm = ChatOpenAI(
-    model="openai/gpt-oss-20b",
-    base_url='https://platform.brojs.ru/jrnl-bh/api/inference/v1',
-    api_key=SecretStr("jrnl_30283ab953615cbb6846ff9940a1eedce0b76d7b2f59a2394f29e74643e6a90d"),
+# ---------- LLM & Embeddings ----------
+from langchain_ollama import Ollama, OllamaEmbeddings
+
+llm = Ollama(
+    model="openai/gpt-oss-20b",  # e.g. "llama3"
+    base_url="http://localhost:11434",
     temperature=0.7,
 )
 
-embeddings = OllamaEmbeddings(model_name="nomic-embed-text")
+embeddings = OllamaEmbeddings(model="nomic-embed-text")
+
+# ---------- Qdrant Vector Store ----------
+from qdrant_client import QdrantClient
+from qdrant_client.http.models import Distance, VectorParams
+from langchain_qdrant import QdrantVectorStore
+
+client = QdrantClient(":memory:")  # in‑memory for demo; replace with path or URL as needed
+# Determine vector size from the embedding model
+sample_vector = embeddings.embed_query("test")[0]
+vector_size = len(sample_vector)
 
-# ---------- Qdrant client ----------
-client = QdrantClient(":memory:")
 client.create_collection(
     collection_name="knowledge_base",
-    vectors_config=VectorParams(size=embeddings.embedding_size, distance=Distance.COSINE),
+    vectors_config=VectorParams(size=vector_size, distance=Distance.COSINE),
 )
-vector_store = QdrantVectorStore(client=client, collection_name="knowledge_base", embedding=embeddings)
 
-# ---------- Text splitter ----------
+vector_store = QdrantVectorStore(
+    client=client,
+    collection_name="knowledge_base",
+    embedding=embeddings,
+)
+
+# ---------- Text Splitter ----------
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+
 splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
 
 # ---------- Tools ----------
-@tool
-def add_to_knowledge_base(content: str, title: str) -> str:
-    """Add a document to the knowledge base."""
-    docs = [Document(page_content=c, metadata={"title": title}) for c in splitter.split_text(content)]
-    vector_store.add_documents(docs)
-    return f"Added {len(docs)} chunks titled '{title}'."
+from langchain.tools import tool
+from langchain_core.documents import Document
 
-@tool
+@tool("search_knowledge_base")
 def search_knowledge_base(query: str, max_results: int = 5) -> str:
-    """Search the knowledge base for relevant information."""
-    results = vector_store.similarity_search_with_score(query, k=max_results)
-    if not results:
-        return "No relevant documents found."
-    reply = ""
-    for i, (doc, score) in enumerate(results, start=1):
-        reply += f"{i}. ({score:.2f}) {doc.metadata.get('title', 'Untitled')}: {doc.page_content[:200]}...\n"
-    return reply.strip()
+    """Search the knowledge base for relevant documents."""
+    docs = vector_store.similarity_search_with_score(query, k=max_results)
+    if not docs:
+        return "No results found."
+    response_lines = []
+    for i, (doc, score) in enumerate(docs, start=1):
+        title = doc.metadata.get("title", "Untitled")
+        snippet = doc.page_content[:200] + ("..." if len(doc.page_content) > 200 else "")
+        response_lines.append(f"{i}. [{score:.2f}] {title}\n{snippet}")
+    return "\n\n".join(response_lines)
+
+@tool("add_to_knowledge_base")
+def add_to_knowledge_base(content: str, title: str = "Untitled") -> str:
+    """Add a new document to the knowledge base."""
+    chunks = splitter.split_text(content)
+    docs = [Document(page_content=c, metadata={"title": title}) for c in chunks]
+    vector_store.add_documents(docs)
+    return f"Added {len(chunks)} chunks under title '{title}'."
 
 # ---------- Agent ----------
-system_prompt = (
-    "You are an assistant that can search and add information to a knowledge base. "
-    "Use the tools `add_to_knowledge_base` and `search_knowledge_base` as needed."
+from langchain.agents import create_agent
+from langchain_core.messages import HumanMessage
+
+agent = create_agent(
+    model=llm,
+    tools=[search_knowledge_base, add_to_knowledge_base],
+    system_message="You are a helpful assistant that can search and update the knowledge base.",
 )
-agent = create_agent(model=llm, tools=[add_to_knowledge_base, search_knowledge_base], system_prompt=system_prompt)
+
+# ---------- Document Loader ----------
+def load_documents_from_dir(directory: str):
+    """Load all .txt files from directory into vector store."""
+    for file_path in Path(directory).glob("*.txt"):
+        text = file_path.read_text(encoding="utf-8")
+        add_to_knowledge_base(text, title=file_path.stem)
 
 # ---------- CLI ----------
 def main():
-    print("RAG Agent CLI. Commands: /add <title> <content>, /search <query>, /quit")
+    print(
+        "Welcome to the RAG Agent. Commands:\n"
+        "/add <title> <file>\n"
+        "/search <query>\n"
+        "/quit\n"
+    )
     while True:
+        user_input = input("> ").strip()
+        if not user_input:
+            continue
+        if user_input.lower() in ("/quit", "exit"):
+            break
+
+        if user_input.startswith("/add"):
             try:
-            inp = input("> ").strip()
-        except EOFError:
-            break
-        if not inp:
-            continue
-        if inp.lower() in ("/quit", "exit"):
-            print("Bye!")
-            break
-        if inp.startswith("/add"):
-            parts = inp.split(maxsplit=2)
-            if len(parts) < 3:
-                print("Usage: /add <title> <content>")
-                continue
-            title, content = parts[1], parts[2]
-            res_msg = agent.invoke({"messages": [{"role": "human", "content": f"/add {title} {content}"}]})
-            for msg in res_msg["messages"]:
-                if hasattr(msg, "tool_calls"):
-                    print(msg.tool_calls[0]["output"])
-        elif inp.startswith("/search"):
-            query = inp[len("/search"):].strip()
+                _, title, file_path = user_input.split(maxsplit=2)
+                content = Path(file_path).read_text(encoding="utf-8")
+                print(add_to_knowledge_base(content, title))
+            except Exception as e:
+                print(f"Error adding document: {e}")
+
+        elif user_input.startswith("/search"):
+            query = user_input[len("/search") :].strip()
             if not query:
-                print("Usage: /search <query>")
+                print("Please provide a search query.")
                 continue
-            res_msg = agent.invoke({"messages": [{"role": "human", "content": f"/search {query}"}]})
-            for msg in res_msg["messages"]:
-                if hasattr(msg, "tool_calls"):
-                    print(msg.tool_calls[0]["output"])
+            result = agent.invoke({"messages": [HumanMessage(content=query)]})
+            for msg in result["messages"]:
+                if hasattr(msg, "content"):
+                    print(msg.content)
+
         else:
-            print("Unknown command. Use /add, /search, or /quit.")
+            # Treat as normal chat message
+            result = agent.invoke({"messages": [HumanMessage(content=user_input)]})
+            for msg in result["messages"]:
+                if hasattr(msg, "content"):
+                    print(msg.content)
 
 if __name__ == "__main__":
+    # Optional: load initial docs from a folder named 'docs'
+    if Path("docs").exists():
+        load_documents_from_dir("docs")
     main()