# Install required packages
!pip install --upgrade --quiet ipywidgets 'pydantic-ai-slim[openrouter,tavily]' python-dotenv braintrust perplexityai

print('✓ Packages installed!')

import os
from getpass import getpass

import braintrust
from dotenv import load_dotenv
from perplexity import Perplexity
from pydantic import BaseModel
from pydantic_ai import Agent, WebSearchTool, WebSearchUserLocation
from pydantic_ai.common_tools.tavily import tavily_search_tool

load_dotenv()

def require_key(name: str, label: str) -> str:
    if not os.getenv(name):
        os.environ[name] = getpass(f"{label}: ")
    return os.environ[name]

require_key("OPENROUTER_API_KEY", "OpenRouter API key")
tavily_key = require_key("TAVILY_API_KEY", "Tavily API key")
require_key("PERPLEXITY_API_KEY", "Perplexity API key")
require_key("BRAINTRUST_API_KEY", "Braintrust API key")

print("Keys initialized.")

from braintrust.wrappers.pydantic_ai import setup_pydantic_ai

if os.getenv("BRAINTRUST_API_KEY"):
    setup_pydantic_ai(project_name="dataharvest-2026")
    print("Braintrust connection established.")
else:
    print("Set BRAINTRUST_API_KEY to enable Braintrust tracing.")

# Model list at https://ai.pydantic.dev/api/models/base/#pydantic_ai.models.KnownModelName
agent = Agent(
    'openrouter:anthropic/claude-haiku-4-5',
    tools=[tavily_search_tool(tavily_key)],
    # Should we add this?
    # builtin_tools=[WebSearchTool(location=WebSearchUserLocation.US)]
)

# Also set up the Perplexity tool
client = Perplexity()

@agent.tool_plain
async def perplexity_research(query: str):
    """Do research on a topic using Perplexity.

    Args:
        query: what to search for
    """
    search = client.search.create(query=query, max_results=10)
    return search.results

# Define the instructions for the agent

instructions = f"""
You are an investigative journalism assistant that triages news tips.

Your task is to:
1. Evaluate the tip for credibility and newsworthiness (score 1-10)
2. Determine if it is valid (score >= 5) or spam/invalid (score < 5)
3. Categorize it into the most appropriate beat from this list: local news, crime, politics, business, health, technology, environment, education, culture, sports, international news
4. Write a concise 2-3 sentence summary
5. Use Perplexity and/or Tavily search to conduct background research on the topic, entities, and location mentioned
6. Research the tipster's name and email to assess credibility (look for public profiles, previous journalism involvement, potential biases)
7. Generate 3-5 specific follow-up questions for further investigation
8. Compose a professional email response:
   - If valid: Thank them, indicate you're reviewing it, and ask the follow-up questions
   - If invalid: Politely decline and thank them for their interest

Be thorough in your research but efficient. Focus on verifiable facts and credible sources. Understand that people don't really know how to write tips, and you should follow up on most of the ones that are not obviously crazy people.
"""

# Define the output from the agent
class ProcessedTip(BaseModel):
    is_valid: bool
    credibility: int
    newsworthiness: int
    validity: int
    beat: str
    summary: str
    research: str
    follow_up_questions: list[str]
    email_response: str

@braintrust.traced
async def triage_tip(name, email, tip):
    prompt = f"""
    Submitter: {name} - {email}
    Tip details: {tip}
    """
    result = await agent.run(
        prompt,
        instructions=instructions,
        output_type=ProcessedTip,
    )
    output = result.output.model_dump()

    return output

name = "Jonathan Soma"
email = "jonathan.soma@gmail.com"
tip = """
The mayor of Fayetteville, NC has been embezzling. I'm the treasurer,
I know it. I have documents, I can prove it.
"""

result = await triage_tip(name, email, tip)

result