Use with PydanticAI

pip install neo4j-agent-memory pydantic-ai

from dataclasses import dataclass, field
from datetime import datetime
import json

from pydantic_ai import Agent, RunContext
from pydantic import BaseModel
from neo4j_agent_memory import MemoryClient


# --- Response Models ---

class ProductRecommendation(BaseModel):
    """Structured product recommendation."""
    products: list[dict]
    reasoning: str
    personalization_used: list[str]


# --- Dependencies ---

@dataclass
class ShoppingDeps:
    memory_client: MemoryClient
    user_id: str
    session_id: str
    current_query: str = ""


# --- Agent Definition ---

shopping_agent = Agent(
    "openai:gpt-4o",
    deps_type=ShoppingDeps,
    result_type=ProductRecommendation,
    system_prompt="""
    You are a personal shopping assistant for an online retailer.

    Your goals:
    1. Understand what the customer is looking for
    2. Use their preferences and history to personalize recommendations
    3. Provide helpful, relevant product suggestions
    4. Remember important details for future interactions

    Always explain why you're recommending specific products based on
    the customer's preferences and past behavior.
    """,
)


# --- Memory Tools ---

@shopping_agent.tool
async def search_products(
    ctx: RunContext[ShoppingDeps],
    query: str,
    category: str | None = None,
    brand: str | None = None,
    max_price: float | None = None,
    limit: int = 10,
) -> str:
    """Search product catalog. Use customer preferences to filter when relevant."""

    # Build property filter
    filters = {}
    if brand:
        filters["brand"] = brand
    if max_price:
        filters["price"] = {"$lte": max_price}

    products = await ctx.deps.memory_client.long_term.search_entities(
        query=query,
        entity_type="PRODUCT",
        property_filter=filters if filters else None,
        limit=limit,
    )

    return json.dumps([
        {
            "name": p.name,
            "brand": p.properties.get("brand"),
            "price": p.properties.get("price"),
            "rating": p.properties.get("rating"),
            "description": p.description,
        }
        for p in products
    ])


@shopping_agent.tool
async def get_customer_preferences(
    ctx: RunContext[ShoppingDeps],
) -> str:
    """Get all stored preferences for the current customer."""

    preferences = await ctx.deps.memory_client.long_term.get_preferences(
        user_id=ctx.deps.user_id,
    )

    # Group by category
    by_category = {}
    for pref in preferences:
        if pref.category not in by_category:
            by_category[pref.category] = []
        by_category[pref.category].append(pref.preference)

    return json.dumps(by_category)


@shopping_agent.tool
async def get_purchase_history(
    ctx: RunContext[ShoppingDeps],
    limit: int = 10,
) -> str:
    """Get customer's recent purchase history from the context graph."""

    # Query the context graph for purchases
    purchases = await ctx.deps.memory_client.long_term.execute_query(
        """
        MATCH (c:Customer {id: $user_id})-[p:PURCHASED]->(product:Product)
        RETURN product.name as name, product.brand as brand,
               product.category as category, p.purchase_date as date
        ORDER BY p.purchase_date DESC
        LIMIT $limit
        """,
        parameters={"user_id": ctx.deps.user_id, "limit": limit},
    )

    return json.dumps(purchases)


@shopping_agent.tool
async def save_preference(
    ctx: RunContext[ShoppingDeps],
    preference: str,
    category: str,
) -> str:
    """Save a new customer preference learned during conversation."""

    await ctx.deps.memory_client.long_term.add_preference(
        user_id=ctx.deps.user_id,
        preference=preference,
        category=category,
        confidence=0.85,
        metadata={
            "source": "conversation",
            "session_id": ctx.deps.session_id,
        },
    )

    return f"Saved preference: {category} - {preference}"


@shopping_agent.tool
async def get_similar_recommendations(
    ctx: RunContext[ShoppingDeps],
) -> str:
    """Find past successful recommendations for similar queries."""

    traces = await ctx.deps.memory_client.reasoning.get_similar_traces(
        task=ctx.deps.current_query,
        user_id=ctx.deps.user_id,
        limit=3,
        success_only=True,
    )

    past_recs = []
    for trace in traces:
        if trace.result and "products" in trace.result:
            past_recs.append({
                "query": trace.task,
                "products_recommended": trace.result["products"][:3],
                "customer_response": trace.result.get("feedback", "unknown"),
            })

    return json.dumps(past_recs)


# --- System Prompt Enhancement ---

@shopping_agent.system_prompt
async def inject_customer_context(ctx: RunContext[ShoppingDeps]) -> str:
    """Inject customer context into system prompt."""

    # Get preferences
    prefs = await ctx.deps.memory_client.long_term.get_preferences(
        user_id=ctx.deps.user_id,
        limit=5,
    )

    if not prefs:
        return ""

    pref_text = "\n".join([f"- {p.category}: {p.preference}" for p in prefs])

    return f"""
## Known Customer Preferences
{pref_text}

Use these preferences to personalize your recommendations.
"""


# --- Main Application ---

async def run_shopping_assistant():
    """Run the shopping assistant with memory."""

    memory_client = MemoryClient(
        neo4j_uri="bolt://localhost:7687",
        neo4j_user="neo4j",
        neo4j_password="password",
    )

    user_id = "CUST-12345"
    session_id = f"shop-{datetime.now().strftime('%Y%m%d%H%M%S')}"

    print("Shopping Assistant Ready! Type 'quit' to exit.\n")

    while True:
        user_input = input("You: ").strip()

        if user_input.lower() == "quit":
            break

        # Store user message
        await memory_client.short_term.add_message(
            role="user",
            content=user_input,
            session_id=session_id,
        )

        # Start reasoning trace
        trace = await memory_client.reasoning.start_trace(
            task=user_input,
            user_id=user_id,
            session_id=session_id,
        )

        deps = ShoppingDeps(
            memory_client=memory_client,
            user_id=user_id,
            session_id=session_id,
            current_query=user_input,
        )

        try:
            result = await shopping_agent.run(user_input, deps=deps)

            # Format response
            response = f"Here are my recommendations:\n\n"
            for i, product in enumerate(result.data.products, 1):
                response += f"{i}. {product['name']} - ${product.get('price', 'N/A')}\n"
            response += f"\n{result.data.reasoning}"

            print(f"\nAssistant: {response}\n")

            # Store assistant response
            await memory_client.short_term.add_message(
                role="assistant",
                content=response,
                session_id=session_id,
            )

            # Complete trace
            await memory_client.reasoning.complete_trace(
                trace_id=trace.id,
                outcome="success",
                result={
                    "products": result.data.products,
                    "personalization": result.data.personalization_used,
                },
            )

        except Exception as e:
            print(f"\nAssistant: Sorry, I encountered an error: {e}\n")

            await memory_client.reasoning.complete_trace(
                trace_id=trace.id,
                outcome="failure",
                error=str(e),
            )


if __name__ == "__main__":
    import asyncio
    asyncio.run(run_shopping_assistant())

from dataclasses import dataclass
from pydantic_ai import Agent, RunContext
from pydantic import BaseModel
from neo4j_agent_memory import MemoryClient


class InvestmentRecommendation(BaseModel):
    """Structured investment recommendation."""
    recommendation: str
    securities: list[dict]
    rationale: str
    risk_considerations: list[str]
    compliance_notes: list[str]


@dataclass
class AdvisoryDeps:
    memory_client: MemoryClient
    client_id: str
    advisor_id: str
    session_id: str


advisory_agent = Agent(
    "openai:gpt-4o",
    deps_type=AdvisoryDeps,
    result_type=InvestmentRecommendation,
    system_prompt="""
    You are a financial advisory assistant helping wealth managers serve their clients.

    Important guidelines:
    1. Always consider the client's risk profile and investment objectives
    2. Reference specific securities with ticker symbols when possible
    3. Note any compliance considerations for recommendations
    4. Use the client's stated preferences and constraints
    5. Document your reasoning for audit purposes

    Never provide specific buy/sell recommendations without noting that final
    decisions should be reviewed by the advisor.
    """,
)


@advisory_agent.tool
async def get_client_profile(ctx: RunContext[AdvisoryDeps]) -> str:
    """Get comprehensive client profile from context graph."""

    # Get client entity with related data
    profile = await ctx.deps.memory_client.long_term.execute_query(
        """
        MATCH (c:Client {id: $client_id})
        OPTIONAL MATCH (c)-[:HAS_ACCOUNT]->(a:Account)
        OPTIONAL MATCH (a)-[:HOLDS]->(p:Position)-[:IN]->(s:Security)
        RETURN c, collect(DISTINCT a) as accounts,
               collect(DISTINCT {security: s, shares: p.shares}) as holdings
        """,
        parameters={"client_id": ctx.deps.client_id},
    )

    # Get investment preferences
    preferences = await ctx.deps.memory_client.long_term.get_preferences(
        user_id=ctx.deps.client_id,
        categories=["risk_profile", "investment_style", "exclusions", "time_horizon"],
    )

    return json.dumps({
        "profile": profile,
        "preferences": [{"category": p.category, "value": p.preference} for p in preferences],
    })


@advisory_agent.tool
async def search_securities(
    ctx: RunContext[AdvisoryDeps],
    query: str,
    sector: str | None = None,
    asset_class: str | None = None,
) -> str:
    """Search for securities matching criteria."""

    filters = {}
    if sector:
        filters["sector"] = sector
    if asset_class:
        filters["asset_class"] = asset_class

    securities = await ctx.deps.memory_client.long_term.search_entities(
        query=query,
        entity_type="SECURITY",
        property_filter=filters if filters else None,
        limit=15,
    )

    return json.dumps([
        {
            "name": s.name,
            "ticker": s.properties.get("ticker"),
            "sector": s.properties.get("sector"),
            "asset_class": s.properties.get("asset_class"),
            "description": s.description,
        }
        for s in securities
    ])


@advisory_agent.tool
async def get_past_recommendations(
    ctx: RunContext[AdvisoryDeps],
    topic: str,
) -> str:
    """Find past successful recommendations for similar situations."""

    traces = await ctx.deps.memory_client.reasoning.get_similar_traces(
        task=topic,
        user_id=ctx.deps.client_id,
        limit=3,
        success_only=True,
    )

    return json.dumps([
        {
            "topic": t.task,
            "recommendation": t.result.get("recommendation"),
            "outcome": t.result.get("client_response"),
        }
        for t in traces
    ])


@advisory_agent.tool
async def record_compliance_note(
    ctx: RunContext[AdvisoryDeps],
    note: str,
    category: str,
) -> str:
    """Record a compliance-relevant note in the context graph."""

    await ctx.deps.memory_client.long_term.add_entity(
        name=f"Compliance Note - {datetime.now().isoformat()}",
        entity_type="COMPLIANCE_NOTE",
        properties={
            "client_id": ctx.deps.client_id,
            "advisor_id": ctx.deps.advisor_id,
            "session_id": ctx.deps.session_id,
            "category": category,
            "note": note,
            "timestamp": datetime.now().isoformat(),
        },
    )

    return f"Recorded compliance note: {category}"


@advisory_agent.system_prompt
async def inject_client_context(ctx: RunContext[AdvisoryDeps]) -> str:
    """Inject client investment profile into system prompt."""

    preferences = await ctx.deps.memory_client.long_term.get_preferences(
        user_id=ctx.deps.client_id,
    )

    if not preferences:
        return ""

    profile_text = "\n".join([f"- {p.category}: {p.preference}" for p in preferences])

    return f"""
## Client Investment Profile
{profile_text}

All recommendations must align with this profile.
"""