This commit is contained in:
+73
-72
@@ -1,93 +1,94 @@
|
||||
"""
|
||||
Core agent implementation.
|
||||
|
||||
This module contains the main Agent class used throughout the project.
|
||||
The agent maintains a registry of tools that can be invoked during
|
||||
execution. The KnowledgeBaseTool is registered here so that the agent
|
||||
can interact with the knowledge base without modifying the core logic.
|
||||
Agent implementation using the latest LangChain API.
|
||||
The agent uses a RetrievalQA chain backed by the Qdrant vector store
|
||||
and the Ollama LLM for local inference.
|
||||
"""
|
||||
|
||||
from typing import Callable, Dict, Any
|
||||
from __future__ import annotations
|
||||
|
||||
# Import the KnowledgeBaseTool but do not alter existing logic
|
||||
from .knowledge_base import KnowledgeBaseTool
|
||||
from typing import Any
|
||||
|
||||
from langchain_ollama import Ollama
|
||||
from langchain.chains import RetrievalQA
|
||||
from langchain.memory import ConversationBufferMemory
|
||||
from langchain.prompts import PromptTemplate
|
||||
|
||||
from .config import OLLAMA_HOST, OLLAMA_PORT, OLLAMA_MODEL
|
||||
from .vector_store import QdrantVectorStore
|
||||
from langchain_core.documents import Document
|
||||
|
||||
|
||||
class Agent:
|
||||
"""
|
||||
A simple agent that can execute registered tools.
|
||||
|
||||
The agent's tool registry maps tool names to callable objects.
|
||||
A simple retrieval-based agent that answers user queries
|
||||
using documents stored in Qdrant and an Ollama LLM.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.tools: Dict[str, Callable[..., Any]] = {}
|
||||
# Register core tools
|
||||
self._register_core_tools()
|
||||
def __init__(self, collection_name: str = "documents"):
|
||||
# Initialize LLM
|
||||
self.llm = Ollama(
|
||||
model=OLLAMA_MODEL,
|
||||
base_url=f"http://{OLLAMA_HOST}:{OLLAMA_PORT}",
|
||||
)
|
||||
|
||||
def _register_core_tools(self) -> None:
|
||||
# Initialize vector store
|
||||
self.vector_store = QdrantVectorStore(collection_name=collection_name)
|
||||
|
||||
# Memory to keep conversation context
|
||||
self.memory = ConversationBufferMemory(memory_key="chat_history")
|
||||
|
||||
# Prompt template
|
||||
self.prompt = PromptTemplate(
|
||||
input_variables=["chat_history", "question"],
|
||||
template=(
|
||||
"You are a helpful assistant. Use the following context to answer the question.\n"
|
||||
"Context:\n{chat_history}\n"
|
||||
"Question: {question}\n"
|
||||
"Answer:"
|
||||
),
|
||||
)
|
||||
|
||||
# RetrievalQA chain
|
||||
self.chain = RetrievalQA.from_chain_type(
|
||||
llm=self.llm,
|
||||
chain_type="stuff",
|
||||
retriever=self.vector_store.as_retriever(),
|
||||
memory=self.memory,
|
||||
return_source_documents=True,
|
||||
chain_type_kwargs={"prompt": self.prompt},
|
||||
)
|
||||
|
||||
def add_documents(self, documents: list[Document]) -> None:
|
||||
"""
|
||||
Register the default set of tools with the agent.
|
||||
Add documents to the underlying vector store.
|
||||
"""
|
||||
# Register the KnowledgeBaseTool under the name 'knowledge_base'
|
||||
self.tools["knowledge_base"] = KnowledgeBaseTool()
|
||||
self.vector_store.add_documents(documents)
|
||||
|
||||
def register_tool(self, name: str, tool: Callable[..., Any]) -> None:
|
||||
def run(self, question: str) -> Any:
|
||||
"""
|
||||
Register a new tool with the agent.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
name : str
|
||||
The name under which the tool will be registered.
|
||||
tool : Callable[..., Any]
|
||||
The tool instance or callable.
|
||||
Run the agent on a user question.
|
||||
Returns the LLM's answer and the source documents.
|
||||
"""
|
||||
self.tools[name] = tool
|
||||
result = self.chain({"question": question})
|
||||
return result
|
||||
|
||||
def run_tool(self, name: str, *args, **kwargs) -> Any:
|
||||
"""
|
||||
Execute a registered tool.
|
||||
def __call__(self, question: str) -> Any:
|
||||
return self.run(question)
|
||||
|
||||
Parameters
|
||||
----------
|
||||
name : str
|
||||
The name of the tool to run.
|
||||
*args, **kwargs
|
||||
Arguments forwarded to the tool.
|
||||
|
||||
Returns
|
||||
-------
|
||||
Any
|
||||
The result of the tool execution.
|
||||
if __name__ == "__main__":
|
||||
import json
|
||||
import sys
|
||||
|
||||
Raises
|
||||
------
|
||||
KeyError
|
||||
If the tool name is not registered.
|
||||
"""
|
||||
if name not in self.tools:
|
||||
raise KeyError(f"Tool '{name}' not found.")
|
||||
tool = self.tools[name]
|
||||
return tool(*args, **kwargs)
|
||||
# Simple CLI usage
|
||||
agent = Agent()
|
||||
if len(sys.argv) > 1:
|
||||
query = " ".join(sys.argv[1:])
|
||||
else:
|
||||
query = input("Enter your question: ")
|
||||
|
||||
# Example method that uses the knowledge base tool
|
||||
def get_fact(self, key: str) -> Any:
|
||||
"""
|
||||
Retrieve a fact from the knowledge base.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
key : str
|
||||
The key to look up.
|
||||
|
||||
Returns
|
||||
-------
|
||||
Any
|
||||
The stored value.
|
||||
"""
|
||||
kb_tool: KnowledgeBaseTool = self.tools["knowledge_base"]
|
||||
return kb_tool.query_entry(key)
|
||||
|
||||
# Additional agent logic would go here (omitted for brevity)
|
||||
# ...
|
||||
response = agent.run(query)
|
||||
print("\nAnswer:\n", response["answer"])
|
||||
print("\nSources:")
|
||||
for doc in response["source_documents"]:
|
||||
print(f"- {doc.metadata.get('source', 'unknown')}")
|
||||
@@ -0,0 +1,26 @@
|
||||
"""
|
||||
Configuration helper for the agent project.
|
||||
Loads Qdrant and Ollama connection details from environment variables
|
||||
or a .env file. Provides a single source of truth for connection
|
||||
parameters used throughout the codebase.
|
||||
"""
|
||||
|
||||
import os
|
||||
from pathlib import Path
|
||||
from dotenv import load_dotenv
|
||||
|
||||
# Load environment variables from .env if present
|
||||
load_dotenv(dotenv_path=Path(__file__).parent.parent / ".env")
|
||||
|
||||
# Qdrant configuration
|
||||
QDRANT_HOST: str = os.getenv("QDRANT_HOST", "localhost")
|
||||
QDRANT_PORT: int = int(os.getenv("QDRANT_PORT", "6333"))
|
||||
QDRANT_API_KEY: str | None = os.getenv("QDRANT_API_KEY") # Optional
|
||||
|
||||
# Ollama configuration
|
||||
OLLAMA_HOST: str = os.getenv("OLLAMA_HOST", "localhost")
|
||||
OLLAMA_PORT: int = int(os.getenv("OLLAMA_PORT", "11434"))
|
||||
OLLAMA_MODEL: str = os.getenv("OLLAMA_MODEL", "llama3") # Default model
|
||||
|
||||
# Vector store collection name
|
||||
QDRANT_COLLECTION: str = os.getenv("QDRANT_COLLECTION", "documents")
|
||||
@@ -0,0 +1,97 @@
|
||||
"""
|
||||
Vector store implementation using Qdrant via langchain-qdrant.
|
||||
Provides a simple interface for adding documents and performing
|
||||
similarity search. Embeddings are generated using OpenAIEmbeddings
|
||||
by default, but can be overridden by passing a custom embedding
|
||||
function.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Iterable, List, Optional
|
||||
|
||||
from langchain.embeddings import OpenAIEmbeddings
|
||||
from langchain_qdrant import Qdrant
|
||||
from langchain.vectorstores import VectorStore
|
||||
from langchain_core.documents import Document
|
||||
|
||||
from .config import (
|
||||
QDRANT_HOST,
|
||||
QDRANT_PORT,
|
||||
QDRANT_API_KEY,
|
||||
QDRANT_COLLECTION,
|
||||
)
|
||||
|
||||
|
||||
class QdrantVectorStore(VectorStore):
|
||||
"""
|
||||
A wrapper around langchain_qdrant.Qdrant that implements the
|
||||
VectorStore interface expected by LangChain chains.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
embeddings: Optional[OpenAIEmbeddings] = None,
|
||||
collection_name: str = QDRANT_COLLECTION,
|
||||
):
|
||||
self.embeddings = embeddings or OpenAIEmbeddings()
|
||||
self.collection_name = collection_name
|
||||
|
||||
# Initialize Qdrant client
|
||||
self.client = Qdrant(
|
||||
host=QDRANT_HOST,
|
||||
port=QDRANT_PORT,
|
||||
api_key=QDRANT_API_KEY,
|
||||
collection_name=self.collection_name,
|
||||
)
|
||||
|
||||
def add_documents(self, documents: Iterable[Document]) -> None:
|
||||
"""
|
||||
Add a collection of documents to the Qdrant store.
|
||||
"""
|
||||
texts = [doc.page_content for doc in documents]
|
||||
metadatas = [doc.metadata for doc in documents]
|
||||
ids = [doc.id for doc in documents if doc.id is not None]
|
||||
|
||||
# Embed the documents
|
||||
embeddings = self.embeddings.embed_documents(texts)
|
||||
|
||||
# Upsert into Qdrant
|
||||
self.client.upsert(
|
||||
embeddings=embeddings,
|
||||
documents=texts,
|
||||
metadatas=metadatas,
|
||||
ids=ids,
|
||||
)
|
||||
|
||||
def similarity_search(
|
||||
self,
|
||||
query: str,
|
||||
k: int = 5,
|
||||
filter: Optional[dict] = None,
|
||||
) -> List[Document]:
|
||||
"""
|
||||
Perform a similarity search against the Qdrant store.
|
||||
"""
|
||||
query_embedding = self.embeddings.embed_query(query)
|
||||
results = self.client.search(
|
||||
query_embedding=query_embedding,
|
||||
limit=k,
|
||||
filter=filter,
|
||||
)
|
||||
# Convert results to Document objects
|
||||
return [
|
||||
Document(
|
||||
page_content=result["payload"]["text"],
|
||||
metadata=result["payload"],
|
||||
id=result["id"],
|
||||
)
|
||||
for result in results
|
||||
]
|
||||
|
||||
# The following methods are required by the VectorStore interface
|
||||
def embed_query(self, query: str) -> List[float]:
|
||||
return self.embeddings.embed_query(query)
|
||||
|
||||
def embed_documents(self, documents: List[str]) -> List[List[float]]:
|
||||
return self.embeddings.embed_documents(documents)
|
||||
Reference in New Issue
Block a user