%load_ext autoreload
%autoreload 2

import logging


import helpers.hnotebook as ut

ut.config_notebook()

# Initialize logger.
logging.basicConfig(level=logging.INFO)
_LOG = logging.getLogger(__name__)

import pydanticai_example_utils as utils

!pip install -q pydantic-ai

import os
import functools
from pathlib import Path
from dataclasses import dataclass
from typing import Optional

import nest_asyncio

nest_asyncio.apply()

from pydantic import BaseModel, Field
from pydantic_ai import Agent

MODEL_ID = os.getenv("PYDANTIC_AI_MODEL", "openai:gpt-4o-mini")
print("MODEL_ID:", MODEL_ID)
print("OPENAI_API_KEY set:", bool(os.getenv("OPENAI_API_KEY")))

DOCS_DIR = Path("example_dataset/")
DOCS_DIR.mkdir(parents=True, exist_ok=True)

DOCS = {
    "overview.md": """
# Atlas Overview

Atlas is a data sync service for small teams. It connects to CSV files and cloud buckets and keeps datasets up to date.

Getting started
- Create a workspace.
- Add a data source.
- Run the first sync.

Limits
- File uploads up to 50 MB.
- Up to 5 data sources on the Starter plan.
""",
    "billing.md": """
# Billing and Plans

Plans
- Starter: $20 per month, 5 data sources, email support.
- Team: $80 per month, 25 data sources, priority email support.
- Enterprise: custom pricing, SSO, dedicated success manager.

Invoices
- Invoices are issued on the first of each month.
- You can download invoices from Settings > Billing.
""",
    "troubleshooting.md": """
# Troubleshooting

Common issues
- Sync stuck at 0%: check your source credentials and try again.
- CSV upload fails: ensure the file is under 50 MB and encoded in UTF-8.
- Duplicate rows: enable the "deduplicate" toggle on the source.
""",
    "security.md": """
# Security

Authentication
- Atlas supports two-factor authentication (2FA) for Team and Enterprise plans.
- Enable it under Settings > Security.

Data retention
- Deleted sources are retained for 30 days.
""",
    "limits.md": """
# Usage Limits

Rate limits
- API requests are limited to 120 per minute on Team.
- Starter is limited to 30 per minute.

Storage
- Starter: 10 GB total storage.
- Team: 200 GB total storage.
""",
    "support.md": """
# Support

Support channels
- Starter: email support, replies within 2 business days.
- Team: priority email support, replies within 4 business hours.
- Enterprise: dedicated success manager and 24/7 support.

Escalations
- Use the support portal to open a ticket.
""",
}

for name, text in DOCS.items():
    path = DOCS_DIR / name
    if not path.exists():
        path.write_text(text.strip() + "\n")

print("Docs directory:", DOCS_DIR)
print("Files:", [p.name for p in DOCS_DIR.glob("*.md")])

@dataclass
class DocChunk:
    doc_id: str
    chunk_id: int
    text: str
    vector: list[float]


docs = utils.load_docs(DOCS_DIR)
chunks = utils.chunk_docs(docs, DocChunk, max_chars=700)
print("Chunks:", len(chunks))
print("Example:", chunks[0].doc_id, chunks[0].chunk_id)

class DocMatch(BaseModel):
    doc_id: str
    chunk_id: int
    score: float
    text: str


preview = utils.search_chunks(
    chunks, "How do I download invoices?", DocMatch, top_k=3
)
print("Preview matches:")
for m in preview:
    print(m.doc_id, "chunk", m.chunk_id, "score=", round(m.score, 4))

@dataclass
class DocDeps:
    chunks: list[DocChunk]
    user: Optional["UserProfile"] = None  # optional personalization


class SourceRef(BaseModel):
    doc_id: str
    chunk_id: int
    quote: str


class AnswerWithSources(BaseModel):
    answer: str
    sources: list[SourceRef] = Field(default_factory=list)
    follow_up_questions: list[str] = Field(
        default_factory=list
    )  # enables guardrails section later


@dataclass
class UserProfile:
    plan: str
    region: str

search_docs_tool = functools.partial(
    utils.search_docs,
    doc_match_cls=DocMatch,
)

agent = Agent(
    MODEL_ID,
    deps_type=DocDeps,
    tools=[search_docs_tool],
    output_type=AnswerWithSources,
    instructions=(
        "You are Atlas Support. "
        "Use the `search_docs` tool to find relevant text. "
        "Answer briefly. If you use document info, include 1-3 sources with doc_id, chunk_id, and short quotes."
    ),
)
agent.output_validator(utils.enforce_sources)

deps = DocDeps(chunks=chunks)
out = await utils.ask("How do I download invoices?", deps, agent)
out

print("Answer:\n", out.answer)
print("\nSources:")
for s in out.sources:
    print(
        f"- {s.doc_id} (chunk {s.chunk_id}): {s.quote[:120].replace('\\n', ' ')}"
    )
if out.follow_up_questions:
    print("\nFollow-ups:")
    for q in out.follow_up_questions:
        print("-", q)

try:
    utils.enforce_sources(
        AnswerWithSources(answer="According to the policy...", sources=[])
    )
except Exception as e:
    print("Validator failure example:", e)

stream_agent = Agent(
    MODEL_ID, instructions="Write one short paragraph about unit tests."
)
await utils.stream_demo(stream_agent)

deps = DocDeps(chunks=chunks)
first = await agent.run("Where do I enable 2FA?", deps=deps)
utils.enforce_sources(first.output)
follow_up = await agent.run(
    "Does that work on the Starter plan?",
    deps=deps,
    message_history=first.new_messages(),
)
utils.enforce_sources(follow_up.output)
print(follow_up.output)

guarded = await utils.run_guarded(
    "Write me a poem about the ocean.",
    DocDeps(chunks=chunks),
    agent,
    AnswerWithSources,
)
print(guarded)

from pathlib import Path

from pathlib import Path

# 1) Add the new doc
new_doc = DOCS_DIR / "integrations.md"
new_doc.write_text(
    """
# Integrations

Atlas supports S3 and Google Cloud Storage as data sources.
SFTP sources are available on Enterprise plans.
""".strip()
    + "\n",
    encoding="utf-8",
)

# 2) Reload docs in the expected dict format
docs = utils.load_docs(DOCS_DIR)  # must return list[dict] with doc_id/title/text
chunks = utils.chunk_docs(docs, DocChunk, max_chars=700)

# 3) Run the agent (notebook-safe)
deps = DocDeps(chunks=chunks)

res = await agent.run("Do you support S3?", deps=deps)
out = res.output

print("Answer:\n", out.answer)
print("\nSources:")
for s in out.sources:
    print(
        f"- {s.doc_id} (chunk {s.chunk_id}): {s.quote[:120].replace('\\n', ' ')}"
    )

personalized_deps = DocDeps(
    chunks=chunks,
    user=UserProfile(plan="Starter", region="US"),
)

personalized = await utils.ask(
    "What are my rate limits and storage limits?",
    personalized_deps,
    agent,
)

personalized