diff --git a/solutions/6a02e23da6fe2e4ac16acf65/solution.py b/solutions/6a02e23da6fe2e4ac16acf65/solution.py index 716e9d4..f291fe7 100644 --- a/solutions/6a02e23da6fe2e4ac16acf65/solution.py +++ b/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 <content>, /search <query>, /quit") + print( + "Welcome to the RAG Agent. Commands:\n" + "/add <title> <file>\n" + "/search <query>\n" + "/quit\n" + ) while True: - try: - inp = input("> ").strip() - except EOFError: - break - if not inp: + user_input = input("> ").strip() + if not user_input: continue - if inp.lower() in ("/quit", "exit"): - print("Bye!") + if user_input.lower() in ("/quit", "exit"): 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() + + if user_input.startswith("/add"): + try: + _, 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() \ No newline at end of file