Building a Knowledge-Augmented Agent

What Is a Knowledge-Augmented Agent?

The Full Retrieval Pipeline

from dataclasses import dataclass, field
from typing import List, Dict, Any

@dataclass
class RetrievalContext:
    question: str
    entities: List[str] = field(default_factory=list)
    graph_context: Dict[str, Any] = field(default_factory=dict)
    vector_documents: List[Dict] = field(default_factory=list)
    combined_context: str = ''
    answer: str = ''
    sources_used: List[str] = field(default_factory=list)

# The agent will populate this object as it works through the pipeline
ctx = RetrievalContext(question='What AI projects is Sam Altman known for?')
print('RetrievalContext created:', ctx.question)

Step 1: Entity Extraction

import spacy

nlp = spacy.load('en_core_web_sm')

def extract_question_entities(question: str) -> List[str]:
    doc = nlp(question)
    entities = list({
        ent.text for ent in doc.ents
        if ent.label_ in ['PERSON', 'ORG', 'GPE', 'PRODUCT', 'WORK_OF_ART']
    })
    return entities

def extract_entities_with_llm_fallback(question: str, client) -> List[str]:
    spacy_entities = extract_question_entities(question)
    
    if spacy_entities:
        return spacy_entities
    
    # Fallback to LLM for questions where spaCy finds nothing
    import json
    response = client.chat.completions.create(
        model='gpt-4o-mini',
        messages=[{
            'role': 'user',
            'content': f'Extract named entities (people, companies, technologies) from: "{question}". Return JSON: {{"entities": ["name1", "name2"]}}'
        }],
        response_format={'type': 'json_object'}
    )
    result = json.loads(response.choices[0].message.content)
    return result.get('entities', [])

question = 'What AI projects is Sam Altman known for?'
entities = extract_question_entities(question)
print('Extracted entities:', entities)

Step 2: Graph Lookup

from neo4j import GraphDatabase

driver = GraphDatabase.driver('bolt://localhost:7687', auth=('neo4j', 'password'))

def get_rich_entity_context(entity_name: str) -> dict:
    with driver.session() as session:
        # Get entity + its relationships
        result = session.run(
            'MATCH (n {name: $name}) '
            'OPTIONAL MATCH (n)-[r]->(target) '
            'RETURN n, labels(n) AS labels, '
            'COLLECT({rel: type(r), target_name: target.name, target_label: labels(target)}) AS outgoing '
            'LIMIT 1',
            name=entity_name
        )
        record = result.single()
        if not record:
            return {'found': False, 'name': entity_name}
        
        return {
            'found': True,
            'name': entity_name,
            'labels': record['labels'],
            'properties': dict(record['n']),
            'connections': [
                c for c in record['outgoing'] if c.get('target_name')
            ][:10]
        }

def format_entity_context_for_llm(entity_ctx: dict) -> str:
    if not entity_ctx.get('found'):
        return f'No knowledge graph data found for "{entity_ctx["name"]}"'
    
    props = entity_ctx.get('properties', {})
    connections = entity_ctx.get('connections', [])
    conn_strs = [f"{c['rel']} -> {c['target_name']}" for c in connections[:5]]
    
    return (
        f"Entity: {entity_ctx['name']} ({', '.join(entity_ctx['labels'])})\n"
        f"Properties: {props}\n"
        f"Relationships: {'; '.join(conn_strs)}"
    )

Step 3: Vector Search

import chromadb
import openai

client = openai.OpenAI(api_key='sk-...')
chroma_client = chromadb.Client()
collection = chroma_client.get_or_create_collection('knowledge_base')

def vector_search(query: str, top_k: int = 5) -> list:
    response = client.embeddings.create(
        model='text-embedding-3-small',
        input=query
    )
    query_embedding = response.data[0].embedding
    
    results = collection.query(
        query_embeddings=[query_embedding],
        n_results=top_k,
        include=['documents', 'metadatas', 'distances']
    )
    
    documents = []
    for i in range(len(results['ids'][0])):
        documents.append({
            'text': results['documents'][0][i],
            'metadata': results['metadatas'][0][i],
            'distance': results['distances'][0][i],
            'relevance': 1 - results['distances'][0][i]  # Convert distance to similarity
        })
    
    return documents

print('Vector search function defined')

Step 4: Combining Context

def combine_context(question: str, graph_contexts: dict, vector_docs: list) -> str:
    sections = []
    
    # Graph facts section
    if graph_contexts:
        graph_parts = ['### Knowledge Graph Facts']
        for entity_name, ctx in graph_contexts.items():
            graph_parts.append(format_entity_context_for_llm(ctx))
        sections.append('\n'.join(graph_parts))
    
    # Vector documents section
    if vector_docs:
        doc_parts = ['### Relevant Documents']
        for i, doc in enumerate(vector_docs[:4]):
            title = doc.get('metadata', {}).get('title', f'Document {i+1}')
            text = doc['text'][:800]  # Limit per document
            relevance = doc.get('relevance', 0)
            doc_parts.append(f'**{title}** (relevance: {relevance:.2f})\n{text}')
        sections.append('\n'.join(doc_parts))
    
    context = '\n\n'.join(sections)
    # Total context budget: ~8000 tokens ~ 32000 chars
    if len(context) > 32000:
        context = context[:32000]
    
    return context

Step 5: LLM Answer Generation

import openai

client = openai.OpenAI(api_key='sk-...')

def generate_answer(question: str, combined_context: str) -> str:
    system_prompt = (
        'You are a knowledgeable assistant. Answer the question using ONLY the provided context. '
        'Cite your sources by mentioning whether a fact came from the knowledge graph or a specific document. '
        'If the context does not contain enough information, say so clearly.'
    )
    
    user_message = (
        f'Context:\n{combined_context}\n\n'
        f'Question: {question}\n\n'
        'Please answer based on the context above.'
    )
    
    response = client.chat.completions.create(
        model='gpt-4o-mini',
        messages=[
            {'role': 'system', 'content': system_prompt},
            {'role': 'user', 'content': user_message}
        ],
        temperature=0.1  # Low temperature for factual answers
    )
    return response.choices[0].message.content

Full Agent Orchestrator

async def knowledge_augmented_agent(question: str) -> RetrievalContext:
    ctx = RetrievalContext(question=question)
    
    # Step 1: Extract entities
    ctx.entities = extract_question_entities(question)
    print(f'Entities: {ctx.entities}')
    
    # Steps 2 & 3: Graph + Vector in parallel
    import asyncio
    from concurrent.futures import ThreadPoolExecutor
    
    executor = ThreadPoolExecutor(max_workers=4)
    loop = asyncio.get_event_loop()
    
    async def graph_step():
        contexts = {}
        for entity in ctx.entities:
            context = await loop.run_in_executor(executor, get_rich_entity_context, entity)
            if context.get('found'):
                contexts[entity] = context
        return contexts
    
    async def vector_step():
        return await loop.run_in_executor(executor, vector_search, question, 5)
    
    ctx.graph_context, ctx.vector_documents = await asyncio.gather(
        graph_step(), vector_step()
    )
    
    # Step 4: Combine
    ctx.combined_context = combine_context(
        question, ctx.graph_context, ctx.vector_documents
    )
    
    # Step 5: Generate answer
    ctx.answer = generate_answer(question, ctx.combined_context)
    
    return ctx

Handling Empty Retrieval

def check_retrieval_quality(graph_contexts: dict, vector_docs: list, threshold: float = 0.7) -> dict:
    has_graph = len(graph_contexts) > 0
    
    # Filter vector docs below relevance threshold
    high_quality_docs = [d for d in vector_docs if d.get('relevance', 0) >= threshold]
    
    return {
        'has_graph_context': has_graph,
        'graph_entity_count': len(graph_contexts),
        'vector_doc_count': len(high_quality_docs),
        'retrieval_quality': 'high' if (has_graph or len(high_quality_docs) >= 2) else 'low',
        'usable_docs': high_quality_docs
    }

def answer_with_fallback(question: str, graph_contexts: dict, vector_docs: list) -> str:
    quality = check_retrieval_quality(graph_contexts, vector_docs)
    
    if quality['retrieval_quality'] == 'low':
        return (
            f'I don\'t have enough information in my knowledge base to answer '
            f'"{question}" confidently. '
            'Please ensure relevant documents are indexed or the knowledge graph '
            'contains the required entities.'
        )
    
    context = combine_context(question, graph_contexts, quality['usable_docs'])
    return generate_answer(question, context)

Caching Retrieval Results

import hashlib
import json
from datetime import datetime, timedelta

class RetrievalCache:
    def __init__(self, ttl_minutes: int = 60):
        self.cache = {}
        self.ttl = timedelta(minutes=ttl_minutes)
    
    def _key(self, namespace: str, value: str) -> str:
        return hashlib.md5(f'{namespace}:{value}'.encode()).hexdigest()
    
    def get(self, namespace: str, value: str):
        key = self._key(namespace, value)
        entry = self.cache.get(key)
        if entry and datetime.now() - entry['ts'] < self.ttl:
            return entry['data']
        return None
    
    def set(self, namespace: str, value: str, data):
        key = self._key(namespace, value)
        self.cache[key] = {'data': data, 'ts': datetime.now()}

cache = RetrievalCache(ttl_minutes=30)

def cached_graph_lookup(entity: str) -> dict:
    cached = cache.get('graph', entity)
    if cached:
        print(f'Cache hit for entity: {entity}')
        return cached
    result = get_rich_entity_context(entity)
    cache.set('graph', entity, result)
    return result

print('Retrieval cache initialized')

Logging and Observability

import logging
import json
from datetime import datetime

logger = logging.getLogger('ka_agent')

def log_agent_run(ctx: 'RetrievalContext', duration_ms: float):
    logger.info(json.dumps({
        'timestamp': datetime.utcnow().isoformat(),
        'question': ctx.question,
        'entities_found': ctx.entities,
        'graph_entities_resolved': list(ctx.graph_context.keys()),
        'vector_docs_retrieved': len(ctx.vector_documents),
        'vector_doc_relevances': [
            round(d.get('relevance', 0), 3)
            for d in ctx.vector_documents
        ],
        'context_length_chars': len(ctx.combined_context),
        'answer_length_chars': len(ctx.answer),
        'duration_ms': round(duration_ms, 1)
    }))

print('Observability logging configured')

Knowledge Check: Knowledge-Augmented Agent

Knowledge-Augmented Agent Summary