Learning Demo: Seed Data

"""
Learning Demo: Seed Data
========================
Runs a few short conversations through the ops assistant so that every
Learning page in AgentOS has data: user profiles, user memories, session
context, entity memories, and decision logs. It also seeds a learned
knowledge insight that one user teaches and another benefits from.

Requires the pgvector container:
    ./cookbook/scripts/run_pgvector.sh

Run:
    .venvs/demo/bin/python cookbook/08_learning/10_demo/seed.py

Then start the AgentOS server with run.py and connect from os.agno.com.
"""

from agents import ops_assistant

ALICE = "alice@vantagelabs.dev"
BEN = "ben@northwind.io"

# (user_id, session_id, message)
CONVERSATIONS = [
    # Alice: profile, preferences, and a session with a clear goal
    (
        ALICE,
        "alice-postgres-upgrade",
        "Hi, I'm Alice Chen, engineering lead at Vantage Labs. "
        "I prefer short, direct answers with code over prose.",
    ),
    (
        ALICE,
        "alice-postgres-upgrade",
        "My goal this week is to upgrade our Postgres cluster from version 15 "
        "to 17 with zero downtime. Help me plan the migration.",
    ),
    (
        ALICE,
        "alice-postgres-upgrade",
        "Some context: Marcus Lee is our infra engineer and owns the Postgres "
        "cluster. The cluster runs on Kubernetes in us-east-1.",
    ),
    (
        ALICE,
        "alice-postgres-upgrade",
        "Should we use logical replication or pg_upgrade for the cutover? "
        "Recommend one and log your decision.",
    ),
    (
        ALICE,
        "alice-postgres-upgrade",
        "Save this for the team: when upgrading Postgres across major "
        "versions, always rehearse the cutover on a clone restored from a "
        "fresh backup before touching production.",
    ),
    # Ben: a second user with different preferences and entities
    (
        BEN,
        "ben-design-system",
        "Hey, I'm Ben Okafor, founder at Northwind. We closed our Series A "
        "round last week. I like detailed answers that walk through trade-offs.",
    ),
    (
        BEN,
        "ben-design-system",
        "We are kicking off the Design System project this quarter and Sarah "
        "Kim will lead it. What should the first milestone be? Pick one and "
        "log your decision.",
    ),
    # Ben benefits from what Alice taught the agent
    (
        BEN,
        "ben-postgres-question",
        "We also need to upgrade Northwind's Postgres soon. Anything the "
        "team has already learned about doing this safely?",
    ),
]

if __name__ == "__main__":
    for user_id, session_id, message in CONVERSATIONS:
        print()
        print("=" * 70)
        print(f"USER: {user_id} | SESSION: {session_id}")
        print("=" * 70)
        ops_assistant.print_response(
            message,
            user_id=user_id,
            session_id=session_id,
            stream=True,
        )

    # ------------------------------------------------------------------
    # Show what the agent learned
    # ------------------------------------------------------------------
    lm = ops_assistant.learning_machine

    print()
    print("=" * 70)
    print("WHAT THE AGENT LEARNED")
    print("=" * 70)

    for user_id in (ALICE, BEN):
        lm.user_profile_store.print(user_id=user_id)
        lm.user_memory_store.print(user_id=user_id)

    lm.session_context_store.print(session_id="alice-postgres-upgrade")
    lm.decision_log_store.print(agent_id="ops-assistant", limit=10)
    lm.learned_knowledge_store.print(query="postgres")

    print()
    print("Entities discovered:")
    seen = set()
    for query in ("postgres", "northwind", "design"):
        for entity in lm.entity_memory_store.search(query=query, limit=5):
            if entity.entity_id not in seen:
                seen.add(entity.entity_id)
                print(f"- {entity.name} ({entity.entity_type})")

    print()
    print("Seed complete. Start the server and explore the Learning pages:")
    print("    .venvs/demo/bin/python cookbook/08_learning/10_demo/run.py")

"""
Learning Demo: Shared Agent
===========================
A single ops assistant with every learning store enabled:

- User Profile: structured fields (name, role, preferences)
- User Memory: unstructured observations about the user
- Session Context: a running summary of each session
- Entity Memory: facts, events, and relationships about external things
- Learned Knowledge: insights that transfer across users (pgvector)
- Decision Log: significant decisions with reasoning

Requires the pgvector container:
    ./cookbook/scripts/run_pgvector.sh
"""

from agno.agent import Agent
from agno.db.postgres import PostgresDb
from agno.knowledge import Knowledge
from agno.knowledge.embedder.openai import OpenAIEmbedder
from agno.learn import (
    LearningMachine,
)
from agno.models.openai import OpenAIResponses
from agno.vectordb.pgvector import PgVector, SearchType

# ---------------------------------------------------------------------------
# Database
# ---------------------------------------------------------------------------
db_url = "postgresql+psycopg://ai:ai@localhost:5532/ai"
db = PostgresDb(id="learning-demo-db", db_url=db_url)

# Learned Knowledge needs a vector store for semantic search.
knowledge = Knowledge(
    vector_db=PgVector(
        db_url=db_url,
        table_name="learning_demo_knowledge",
        search_type=SearchType.hybrid,
        embedder=OpenAIEmbedder(id="text-embedding-3-small"),
    ),
)

# ---------------------------------------------------------------------------
# Learning Machine: all six stores enabled
# ---------------------------------------------------------------------------
learning = LearningMachine(
    db=db,
    model=OpenAIResponses(id="gpt-5.5"),
    knowledge=knowledge,
    user_profile=True,
    user_memory=True,
    session_context=True,
    entity_memory=True,
    learned_knowledge=True,
    decision_log=True,
)

# ---------------------------------------------------------------------------
# Agent
# ---------------------------------------------------------------------------
ops_assistant = Agent(
    id="ops-assistant",
    name="Ops Assistant",
    model=OpenAIResponses(id="gpt-5.5"),
    db=db,
    learning=learning,
    instructions=[
        "You are an engineering operations assistant.",
        "Keep answers short and practical.",
        "Search your learnings before answering substantive questions.",
        "When the user shares a team-wide insight or asks you to remember one, save it with the save_learning tool.",
        "When you make a significant recommendation, record it with the log_decision tool, including your reasoning and the alternatives you considered.",
    ],
    markdown=True,
)