diff --git a/solutions/6a02e23da6fe2e4ac16acf65/solution.py b/solutions/6a02e23da6fe2e4ac16acf65/solution.py new file mode 100644 index 0000000..d9f09cc --- /dev/null +++ b/solutions/6a02e23da6fe2e4ac16acf65/solution.py @@ -0,0 +1,90 @@ +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_qdrant import QdrantVectorStore +from qdrant_client import QdrantClient +from qdrant_client.http.models import Distance, VectorParams +from langchain_core.messages import HumanMessage + +# ---------- Qdrant setup ---------- +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, +) + +# ---------- 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) + ) + +@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." + +# ---------- 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.", +) + +# ---------- CLI ---------- +def main(): + print("RAG Agent CLI. Commands: /add