This commit is contained in:
+82
-60
@@ -1,75 +1,97 @@
|
||||
"""
|
||||
Agent implementation that uses LangChain to interact with the knowledge base.
|
||||
RAG Agent implementation.
|
||||
|
||||
Stores documents in memory, retrieves top-k relevant documents using cosine similarity,
|
||||
constructs a prompt with context, and generates a response via Ollama LLM.
|
||||
"""
|
||||
|
||||
import os
|
||||
from typing import List
|
||||
from langchain import OpenAI
|
||||
from langchain.agents import ZeroShotAgent, AgentExecutor
|
||||
from langchain.tools import BaseTool
|
||||
from knowledge_base import add_to_knowledge_base, search_knowledge_base
|
||||
from langchain.schema import AgentAction, AgentFinish
|
||||
from typing import List, Tuple
|
||||
import numpy as np
|
||||
|
||||
class SearchTool(BaseTool):
|
||||
from .embeddings import embed, cosine_similarity
|
||||
from .llm import chat
|
||||
|
||||
|
||||
class RAGAgent:
|
||||
"""
|
||||
Tool to search the knowledge base.
|
||||
Retrieval-Augmented Generation Agent.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
model : str, optional
|
||||
Ollama model to use for embeddings and LLM. Defaults to "llama2".
|
||||
top_k : int, optional
|
||||
Number of top documents to retrieve. Defaults to 3.
|
||||
"""
|
||||
name = "search_knowledge_base"
|
||||
description = (
|
||||
"Search the knowledge base for relevant information. "
|
||||
"Input: query string."
|
||||
)
|
||||
|
||||
def _run(self, query: str) -> str:
|
||||
docs: List = search_knowledge_base(query)
|
||||
if not docs:
|
||||
return "No relevant documents found."
|
||||
return "\n---\n".join([doc.page_content for doc in docs])
|
||||
def __init__(self, model: str = "llama2", top_k: int = 3):
|
||||
self.model = model
|
||||
self.top_k = top_k
|
||||
# Store tuples of (embedding, text)
|
||||
self._store: List[Tuple[List[float], str]] = []
|
||||
|
||||
class AddTool(BaseTool):
|
||||
"""
|
||||
Tool to add new knowledge to the knowledge base.
|
||||
"""
|
||||
name = "add_to_knowledge_base"
|
||||
description = (
|
||||
"Add new knowledge to the knowledge base. "
|
||||
"Input: content string."
|
||||
)
|
||||
def add_document(self, text: str) -> None:
|
||||
"""
|
||||
Add a document to the in-memory vector store.
|
||||
|
||||
def _run(self, content: str) -> str:
|
||||
return add_to_knowledge_base(content)
|
||||
Parameters
|
||||
----------
|
||||
text : str
|
||||
Document text.
|
||||
"""
|
||||
vec = embed(text, model=self.model)
|
||||
self._store.append((vec, text))
|
||||
|
||||
# Instantiate tools
|
||||
tools = [SearchTool(), AddTool()]
|
||||
def _retrieve(self, query: str) -> List[str]:
|
||||
"""
|
||||
Retrieve top-k documents relevant to the query.
|
||||
|
||||
# LLM configuration
|
||||
llm = OpenAI(temperature=0)
|
||||
Parameters
|
||||
----------
|
||||
query : str
|
||||
Query text.
|
||||
|
||||
# Create the agent
|
||||
agent = ZeroShotAgent(llm=llm, tools=tools)
|
||||
Returns
|
||||
-------
|
||||
List[str]
|
||||
List of retrieved document texts.
|
||||
"""
|
||||
query_vec = embed(query, model=self.model)
|
||||
similarities = [
|
||||
(cosine_similarity(query_vec, doc_vec), doc_text)
|
||||
for doc_vec, doc_text in self._store
|
||||
]
|
||||
# Sort by similarity descending
|
||||
similarities.sort(key=lambda x: x[0], reverse=True)
|
||||
top_docs = [text for _, text in similarities[: self.top_k]]
|
||||
return top_docs
|
||||
|
||||
# Executor that runs the agent
|
||||
agent_executor = AgentExecutor.from_agent_and_tools(
|
||||
agent=agent,
|
||||
tools=tools,
|
||||
verbose=True
|
||||
)
|
||||
def get_response(self, query: str) -> str:
|
||||
"""
|
||||
Generate a response to the query using retrieved context.
|
||||
|
||||
def main() -> None:
|
||||
"""
|
||||
Simple REPL to interact with the agent.
|
||||
"""
|
||||
print("Welcome to the RAG Agent. Type 'exit' to quit.")
|
||||
while True:
|
||||
user_input = input("User: ")
|
||||
if user_input.lower() in ("exit", "quit"):
|
||||
print("Goodbye!")
|
||||
break
|
||||
try:
|
||||
response = agent_executor.run(user_input)
|
||||
print(f"Agent: {response}")
|
||||
except Exception as e:
|
||||
print(f"Error: {e}")
|
||||
Parameters
|
||||
----------
|
||||
query : str
|
||||
User query.
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Returns
|
||||
-------
|
||||
str
|
||||
Generated answer.
|
||||
"""
|
||||
context_docs = self._retrieve(query)
|
||||
context = "\n\n".join(context_docs)
|
||||
|
||||
system_prompt = (
|
||||
"You are an assistant that uses the provided context to answer the question."
|
||||
)
|
||||
user_prompt = f"Context:\n{context}\n\nQuestion: {query}\nAnswer:"
|
||||
|
||||
messages = [
|
||||
{"role": "system", "content": system_prompt},
|
||||
{"role": "user", "content": user_prompt},
|
||||
]
|
||||
|
||||
response = chat(messages, model=self.model)
|
||||
return response.strip()
|
||||
+68
-11
@@ -1,18 +1,75 @@
|
||||
"""
|
||||
Embeddings module that provides an OllamaEmbeddings instance.
|
||||
Embeddings module using Ollama.
|
||||
|
||||
Provides a simple caching layer and a function to embed text using Ollama's
|
||||
embedding endpoint. No OpenAI services are used.
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
from langchain_community.embeddings import OllamaEmbeddings
|
||||
from typing import List, Dict
|
||||
|
||||
def get_embedding_model() -> OllamaEmbeddings:
|
||||
"""
|
||||
Returns an OllamaEmbeddings instance configured with the model name
|
||||
specified by the OLLAMA_MODEL environment variable or defaults to
|
||||
'mistral'.
|
||||
import ollama
|
||||
import numpy as np
|
||||
|
||||
Returns:
|
||||
OllamaEmbeddings: The embedding model instance.
|
||||
# Cache to avoid repeated calls for the same text
|
||||
_EMBED_CACHE: Dict[str, List[float]] = {}
|
||||
|
||||
|
||||
def embed(text: str, model: str = "llama2") -> List[float]:
|
||||
"""
|
||||
model_name = os.getenv("OLLAMA_MODEL", "mistral")
|
||||
return OllamaEmbeddings(model=model_name)
|
||||
Generate an embedding vector for the given text using Ollama.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
text : str
|
||||
The text to embed.
|
||||
model : str, optional
|
||||
The Ollama model to use for embeddings. Defaults to "llama2".
|
||||
|
||||
Returns
|
||||
-------
|
||||
List[float]
|
||||
The embedding vector.
|
||||
"""
|
||||
if text in _EMBED_CACHE:
|
||||
return _EMBED_CACHE[text]
|
||||
|
||||
# Ollama expects a dict with "model" and "prompt"
|
||||
payload = {"model": model, "prompt": text}
|
||||
try:
|
||||
response = ollama.embeddings(payload)
|
||||
except Exception as exc:
|
||||
raise RuntimeError(f"Failed to get embeddings from Ollama: {exc}") from exc
|
||||
|
||||
# Ollama returns a dict with "embedding" key
|
||||
embedding = response.get("embedding")
|
||||
if embedding is None:
|
||||
raise ValueError("Ollama response missing 'embedding' field")
|
||||
|
||||
_EMBED_CACHE[text] = embedding
|
||||
return embedding
|
||||
|
||||
|
||||
def cosine_similarity(vec1: List[float], vec2: List[float]) -> float:
|
||||
"""
|
||||
Compute cosine similarity between two vectors.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
vec1, vec2 : List[float]
|
||||
Input vectors.
|
||||
|
||||
Returns
|
||||
-------
|
||||
float
|
||||
Cosine similarity score.
|
||||
"""
|
||||
v1 = np.array(vec1)
|
||||
v2 = np.array(vec2)
|
||||
dot = np.dot(v1, v2)
|
||||
norm1 = np.linalg.norm(v1)
|
||||
norm2 = np.linalg.norm(v2)
|
||||
if norm1 == 0 or norm2 == 0:
|
||||
return 0.0
|
||||
return dot / (norm1 * norm2)
|
||||
+63
@@ -0,0 +1,63 @@
|
||||
"""
|
||||
LLM inference module using Ollama.
|
||||
|
||||
Provides a simple wrapper around Ollama's chat endpoint.
|
||||
"""
|
||||
|
||||
import json
|
||||
from typing import List, Dict, Any, Optional
|
||||
|
||||
import ollama
|
||||
|
||||
|
||||
def chat(
|
||||
messages: List[Dict[str, str]],
|
||||
model: str = "llama2",
|
||||
stream: bool = False,
|
||||
temperature: float = 0.7,
|
||||
max_tokens: Optional[int] = None,
|
||||
) -> str:
|
||||
"""
|
||||
Generate a response from the LLM using Ollama.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
messages : List[Dict[str, str]]
|
||||
List of messages in the format expected by Ollama chat API.
|
||||
model : str, optional
|
||||
The Ollama model to use. Defaults to "llama2".
|
||||
stream : bool, optional
|
||||
Whether to stream the response. Defaults to False.
|
||||
temperature : float, optional
|
||||
Sampling temperature. Defaults to 0.7.
|
||||
max_tokens : int, optional
|
||||
Maximum number of tokens to generate.
|
||||
|
||||
Returns
|
||||
-------
|
||||
str
|
||||
The generated response text.
|
||||
"""
|
||||
payload = {
|
||||
"model": model,
|
||||
"messages": messages,
|
||||
"temperature": temperature,
|
||||
}
|
||||
if max_tokens is not None:
|
||||
payload["max_tokens"] = max_tokens
|
||||
|
||||
try:
|
||||
if stream:
|
||||
# Streaming returns a generator of dicts
|
||||
response_gen = ollama.chat(payload, stream=True)
|
||||
chunks = []
|
||||
for chunk in response_gen:
|
||||
# Each chunk contains a "message" dict with "content"
|
||||
content = chunk.get("message", {}).get("content", "")
|
||||
chunks.append(content)
|
||||
return "".join(chunks)
|
||||
else:
|
||||
response = ollama.chat(payload)
|
||||
return response.get("message", {}).get("content", "")
|
||||
except Exception as exc:
|
||||
raise RuntimeError(f"Failed to get chat response from Ollama: {exc}") from exc
|
||||
+35
-12
@@ -1,21 +1,44 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Main entry point for the knowledge‑base agent.
|
||||
FastAPI application exposing the RAG agent as a REST endpoint.
|
||||
|
||||
This file is optional but useful for running the agent in a container.
|
||||
"""
|
||||
|
||||
from .knowledge_base import KnowledgeBase
|
||||
from .tools.knowledge_base_tool import KnowledgeBaseTool
|
||||
from .cli import run_cli
|
||||
from fastapi import FastAPI, HTTPException
|
||||
from pydantic import BaseModel
|
||||
|
||||
from .agent import RAGAgent
|
||||
|
||||
app = FastAPI(title="RAG Agent API")
|
||||
|
||||
# Initialize a global agent instance
|
||||
agent = RAGAgent(model="llama2", top_k=3)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
class Document(BaseModel):
|
||||
text: str
|
||||
|
||||
|
||||
class Query(BaseModel):
|
||||
query: str
|
||||
|
||||
|
||||
@app.post("/documents")
|
||||
def add_document(doc: Document):
|
||||
"""
|
||||
Create the knowledge base, wrap it in a tool, and start the CLI.
|
||||
Add a document to the agent's memory.
|
||||
"""
|
||||
kb = KnowledgeBase()
|
||||
kb_tool = KnowledgeBaseTool(kb)
|
||||
run_cli(kb_tool)
|
||||
agent.add_document(doc.text)
|
||||
return {"status": "added"}
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@app.post("/ask")
|
||||
def ask(query: Query):
|
||||
"""
|
||||
Get an answer to a query using the RAG agent.
|
||||
"""
|
||||
try:
|
||||
answer = agent.get_response(query.query)
|
||||
except Exception as exc:
|
||||
raise HTTPException(status_code=500, detail=str(exc))
|
||||
return {"answer": answer}
|
||||
Reference in New Issue
Block a user