diff --git a/solutions/6a02e23da6fe2e4ac16acf65/solution.py b/solutions/6a02e23da6fe2e4ac16acf65/solution.py index d9f09cc..716e9d4 100644 --- a/solutions/6a02e23da6fe2e4ac16acf65/solution.py +++ b/solutions/6a02e23da6fe2e4ac16acf65/solution.py @@ -1,58 +1,64 @@ -from pathlib import Path - -# LLM and embeddings via Ollama -from langchain_ollama import ChatOllama, OllamaEmbeddings -from langchain.tools import tool -from langchain.agents import create_agent -from langchain_core.documents import Document +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.messages import HumanMessage +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 -# ---------- Qdrant setup ---------- +# ---------- 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"), + temperature=0.7, +) + +embeddings = OllamaEmbeddings(model_name="nomic-embed-text") + +# ---------- Qdrant client ---------- client = QdrantClient(":memory:") client.create_collection( - collection_name="knowledge", - vectors_config=VectorParams(size=1024, distance=Distance.COSINE), -) -embeddings = OllamaEmbeddings(model="nomic-embed-text") -vector_store = QdrantVectorStore( - client=client, - collection_name="knowledge", - embedding=embeddings, + collection_name="knowledge_base", + vectors_config=VectorParams(size=embeddings.embedding_size, distance=Distance.COSINE), ) +vector_store = QdrantVectorStore(client=client, collection_name="knowledge_base", embedding=embeddings) + +# ---------- Text splitter ---------- +splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50) # ---------- Tools ---------- @tool -def search_knowledge_base(query: str, max_results: int = 5) -> str: - """Search the knowledge base for relevant documents.""" - docs_with_score = vector_store.similarity_search_with_score(query, k=max_results) - if not docs_with_score: - return "No results found." - return "\n".join( - f"{i+1}. {doc.page_content[:200]}..." - for i, (doc, _) in enumerate(docs_with_score) - ) +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}'." @tool -def add_to_knowledge_base(content: str, title: str = "") -> str: - """Add a new document to the knowledge base.""" - doc = Document(page_content=content, metadata={"title": title}) - vector_store.add_documents([doc]) - return f"Document '{title}' added." +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() # ---------- Agent ---------- -llm = ChatOllama(model="llama3") -agent = create_agent( - model=llm, - tools=[search_knowledge_base, add_to_knowledge_base], - system_prompt="You are a helpful assistant that can search and store knowledge.", +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." ) +agent = create_agent(model=llm, tools=[add_to_knowledge_base, search_knowledge_base], system_prompt=system_prompt) # ---------- CLI ---------- def main(): - print("RAG Agent CLI. Commands: /add | <content>, /search <query>, /quit") + print("RAG Agent CLI. Commands: /add <title> <content>, /search <query>, /quit") while True: try: inp = input("> ").strip() @@ -60,31 +66,30 @@ def main(): break if not inp: continue - if inp.lower() in ("quit", "/quit"): + if inp.lower() in ("/quit", "exit"): print("Bye!") break - - # Add document - if inp.startswith("/add "): - _, rest = inp.split(maxsplit=1) - try: - title, content = rest.split("|", 1) - except ValueError: - print("Usage: /add <title> | <content>") + if inp.startswith("/add"): + parts = inp.split(maxsplit=2) + if len(parts) < 3: + print("Usage: /add <title> <content>") continue - res = agent.invoke({"messages": [HumanMessage(content=f"Add document {title}")]}) - print(res.messages[-1].content) - - # Search documents - elif inp.startswith("/search "): - query = inp[len("/search "):] - res = agent.invoke({"messages": [HumanMessage(content=f"Search for {query}")]}) - print(res.messages[-1].content) - - # General chat + 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 not query: + print("Usage: /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"]) else: - res = agent.invoke({"messages": [HumanMessage(content=inp)]}) - print(res.messages[-1].content) + print("Unknown command. Use /add, /search, or /quit.") if __name__ == "__main__": main() \ No newline at end of file