Use with LlamaIndex

pip install neo4j-agent-memory llama-index llama-index-llms-openai

import json
from datetime import datetime
from llama_index.core import (
    VectorStoreIndex,
    SimpleDirectoryReader,
    Settings,
)
from llama_index.core.tools import QueryEngineTool, ToolMetadata
from llama_index.core.agent import ReActAgent
from llama_index.llms.openai import OpenAI
from llama_index.core.memory import ChatMemoryBuffer
from neo4j_agent_memory import MemoryClient


# --- Initialize ---

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

Settings.llm = OpenAI(model="gpt-4o", temperature=0)


# --- Create Document Indexes ---

# Index for regulatory documents
regulatory_docs = SimpleDirectoryReader("./data/regulatory").load_data()
regulatory_index = VectorStoreIndex.from_documents(regulatory_docs)

# Index for market research
research_docs = SimpleDirectoryReader("./data/research").load_data()
research_index = VectorStoreIndex.from_documents(research_docs)


# --- Create Context Graph Tools ---

from llama_index.core.tools import FunctionTool


def get_client_profile(client_id: str) -> str:
    """Get comprehensive client profile from context graph."""
    import asyncio
    loop = asyncio.get_event_loop()

    # Get preferences
    preferences = loop.run_until_complete(
        memory_client.long_term.get_preferences(
            user_id=client_id,
        )
    )

    # Get related entities (holdings, accounts)
    holdings = loop.run_until_complete(
        memory_client.long_term.execute_query(
            """
            MATCH (c:Client {id: $client_id})-[:HAS_ACCOUNT]->(a:Account)
                  -[:HOLDS]->(p:Position)-[:IN]->(s:Security)
            RETURN a.type as account_type, s.name as security,
                   s.ticker as ticker, p.shares as shares
            """,
            parameters={"client_id": client_id},
        )
    )

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


def search_securities(query: str, sector: str = None) -> str:
    """Search for securities in the context graph."""
    import asyncio
    loop = asyncio.get_event_loop()

    filters = {"sector": sector} if sector else None

    securities = loop.run_until_complete(
        memory_client.long_term.search_entities(
            query=query,
            entity_type="SECURITY",
            property_filter=filters,
            limit=10,
        )
    )

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


def get_past_recommendations(client_id: str, topic: str) -> str:
    """Get past successful recommendations for similar topics."""
    import asyncio
    loop = asyncio.get_event_loop()

    traces = loop.run_until_complete(
        memory_client.reasoning.get_similar_traces(
            task=topic,
            user_id=client_id,
            limit=3,
            success_only=True,
        )
    )

    return json.dumps([
        {
            "topic": t.task,
            "recommendation": t.result.get("recommendation") if t.result else None,
            "client_feedback": t.result.get("feedback") if t.result else None,
        }
        for t in traces
    ])


# --- Build Tools ---

tools = [
    # Document query tools
    QueryEngineTool(
        query_engine=regulatory_index.as_query_engine(),
        metadata=ToolMetadata(
            name="regulatory_search",
            description="Search regulatory documents for compliance and rules information.",
        ),
    ),
    QueryEngineTool(
        query_engine=research_index.as_query_engine(),
        metadata=ToolMetadata(
            name="research_search",
            description="Search market research reports for analysis and insights.",
        ),
    ),
    # Context graph tools
    FunctionTool.from_defaults(
        fn=get_client_profile,
        name="get_client_profile",
        description="Get client's investment profile, preferences, and current holdings.",
    ),
    FunctionTool.from_defaults(
        fn=search_securities,
        name="search_securities",
        description="Search for securities by description or characteristics.",
    ),
    FunctionTool.from_defaults(
        fn=get_past_recommendations,
        name="get_past_recommendations",
        description="Find past successful recommendations for similar client situations.",
    ),
]


# --- Create Agent with Memory ---

def create_financial_agent(client_id: str, session_id: str) -> ReActAgent:
    """Create financial advisory agent with context graph memory."""

    # Chat store for message persistence
    chat_store = Neo4jChatStore(memory_client)

    memory = ChatMemoryBuffer.from_defaults(
        chat_store=chat_store,
        chat_store_key=session_id,
        token_limit=4000,
    )

    agent = ReActAgent.from_tools(
        tools=tools,
        memory=memory,
        verbose=True,
        system_prompt=f"""You are a financial advisory assistant helping wealth managers.

Current client: {client_id}

Your capabilities:
1. Search regulatory documents for compliance guidance
2. Search market research for investment insights
3. Access client profiles, preferences, and holdings
4. Search for securities matching specific criteria
5. Review past successful recommendations

Guidelines:
- Always consider the client's risk profile and preferences
- Reference specific securities with ticker symbols
- Note compliance considerations when relevant
- Learn from past successful recommendations

Remember: All recommendations should be reviewed by the advisor before execution.""",
    )

    return agent


# --- Main Application ---

async def run_financial_assistant():
    client_id = "CL-78901"
    session_id = f"advisory-{datetime.now().strftime('%Y%m%d%H%M%S')}"

    agent = create_financial_agent(client_id, session_id)

    print("Financial Advisory Assistant (type 'quit' to exit)\n")

    while True:
        user_input = input("Advisor: ").strip()
        if user_input.lower() == "quit":
            break

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

        try:
            response = agent.chat(user_input)
            print(f"\nAssistant: {response}\n")

            await memory_client.reasoning.complete_trace(
                trace_id=trace.id,
                outcome="success",
                result={"response": str(response)[:500]},
            )
        except Exception as e:
            print(f"\nError: {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_financial_assistant())

from llama_index.core import VectorStoreIndex, Document
from llama_index.core.tools import FunctionTool, QueryEngineTool
from llama_index.core.agent import ReActAgent
from llama_index.llms.openai import OpenAI


def create_product_index(memory_client: MemoryClient) -> VectorStoreIndex:
    """Create product index from context graph."""
    import asyncio
    loop = asyncio.get_event_loop()

    # Get all products from context graph
    products = loop.run_until_complete(
        memory_client.long_term.search_entities(
            query="",
            entity_type="PRODUCT",
            limit=1000,
        )
    )

    # Convert to LlamaIndex documents
    documents = []
    for p in products:
        doc_text = f"""
Product: {p.name}
Brand: {p.properties.get('brand', 'Unknown')}
Category: {p.properties.get('category', 'Unknown')}
Price: ${p.properties.get('price', 'N/A')}
Rating: {p.properties.get('rating', 'N/A')}/5
Description: {p.description or 'No description available'}
        """
        documents.append(Document(
            text=doc_text,
            metadata={
                "product_name": p.name,
                "brand": p.properties.get("brand"),
                "category": p.properties.get("category"),
            },
        ))

    return VectorStoreIndex.from_documents(documents)


def get_customer_preferences(customer_id: str) -> str:
    """Get customer preferences from context graph."""
    import asyncio
    loop = asyncio.get_event_loop()

    preferences = loop.run_until_complete(
        memory_client.long_term.get_preferences(user_id=customer_id)
    )

    by_category = {}
    for p in preferences:
        if p.category not in by_category:
            by_category[p.category] = []
        by_category[p.category].append(p.preference)

    return json.dumps(by_category)


def get_purchase_history(customer_id: str) -> str:
    """Get customer purchase history."""
    import asyncio
    loop = asyncio.get_event_loop()

    purchases = loop.run_until_complete(
        memory_client.long_term.execute_query(
            """
            MATCH (c:Customer {id: $customer_id})-[p:PURCHASED]->(product:Product)
            RETURN product.name as product, product.brand as brand,
                   p.purchase_date as date, p.rating as rating
            ORDER BY p.purchase_date DESC
            LIMIT 10
            """,
            parameters={"customer_id": customer_id},
        )
    )

    return json.dumps(purchases)


def save_preference(customer_id: str, preference: str, category: str) -> str:
    """Save a customer preference."""
    import asyncio
    loop = asyncio.get_event_loop()

    loop.run_until_complete(
        memory_client.long_term.add_preference(
            user_id=customer_id,
            preference=preference,
            category=category,
        )
    )

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


def create_shopping_agent(
    memory_client: MemoryClient,
    customer_id: str,
    session_id: str,
) -> ReActAgent:
    """Create shopping assistant with product knowledge and customer context."""

    # Create product index
    product_index = create_product_index(memory_client)

    # Create tools
    tools = [
        QueryEngineTool(
            query_engine=product_index.as_query_engine(similarity_top_k=10),
            metadata=ToolMetadata(
                name="search_products",
                description="Search product catalog by description, category, or features.",
            ),
        ),
        FunctionTool.from_defaults(
            fn=lambda: get_customer_preferences(customer_id),
            name="get_preferences",
            description="Get current customer's stored preferences.",
        ),
        FunctionTool.from_defaults(
            fn=lambda: get_purchase_history(customer_id),
            name="get_purchase_history",
            description="Get customer's recent purchase history.",
        ),
        FunctionTool.from_defaults(
            fn=lambda pref, cat: save_preference(customer_id, pref, cat),
            name="save_preference",
            description="Save a new customer preference. Parameters: preference (str), category (str).",
        ),
    ]

    # Create chat store
    chat_store = Neo4jChatStore(memory_client)
    memory = ChatMemoryBuffer.from_defaults(
        chat_store=chat_store,
        chat_store_key=session_id,
        token_limit=3000,
    )

    return ReActAgent.from_tools(
        tools=tools,
        memory=memory,
        verbose=True,
        llm=OpenAI(model="gpt-4o"),
        system_prompt=f"""You are a personal shopping assistant.

Customer: {customer_id}

Your capabilities:
1. Search the product catalog for items matching requests
2. Access customer preferences for personalization
3. Review purchase history for context
4. Save new preferences when customers express them

Always:
- Check customer preferences before making recommendations
- Personalize suggestions based on past purchases
- Save important preferences for future interactions
- Explain why you're recommending specific products""",
    )