Spaces:
Runtime error
Runtime error
File size: 3,012 Bytes
cf0f589 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 |
from dotenv import load_dotenv
from dataclasses import dataclass
from datetime import datetime
from typing import Any
import pathlib
import asyncio
import sys
import os
import logfire
from httpx import AsyncClient
from pydantic_ai import Agent
from openai import AsyncOpenAI, OpenAI
from pydantic_ai.models.openai import OpenAIModel
load_dotenv()
def get_model():
llm = os.getenv('MODEL_CHOICE', 'gpt-4o-mini')
base_url = os.getenv('BASE_URL', 'https://api.openai.com/v1')
api_key = os.getenv('LLM_API_KEY', 'no-api-key-provided')
return OpenAIModel(
llm,
base_url=base_url,
api_key=api_key
)
@dataclass
class Deps:
client: AsyncClient
brave_api_key: str | None
agent = Agent(model=get_model(), deps_type=Deps)
@agent.tool
async def search_web(
ctx: RunContext[Deps], web_query: str
) -> str:
"""Search the web given a query defined to answer the user's question.
Args:
ctx: The context.
web_query: The query for the web search.
Returns:
str: The search results as a formatted string.
"""
if ctx.deps.brave_api_key is None:
return "This is a test web search result. Please provide a Brave API key to get real search results."
headers = {
'X-Subscription-Token': ctx.deps.brave_api_key,
'Accept': 'application/json',
}
with logfire.span('calling Brave search API', query=web_query) as span:
r = await ctx.deps.client.get(
'https://api.search.brave.com/res/v1/web/search',
params={
'q': web_query,
'count': 5,
'text_decorations': True,
'search_lang': 'en'
},
headers=headers
)
r.raise_for_status()
data = r.json()
span.set_attribute('response', data)
results = []
# Add web results in a nice formatted way
web_results = data.get('web', {}).get('results', [])
for item in web_results[:3]:
title = item.get('title', '')
description = item.get('description', '')
url = item.get('url', '')
if title and description:
results.append(f"Title: {title}\nSummary: {description}\nSource: {url}\n")
return "\n".join(results) if results else "No results found for the query."
async def main():
while True:
# Example: Search the web to find the newest local LLMs.
user_input = input("\n[You] ")
# Check if user wants to exit
if user_input.lower() in ['exit', 'quit', 'bye', 'goodbye']:
print("Goodbye!")
break
async with AsyncClient() as client:
brave_api_key = os.getenv('BRAVE_API_KEY', None)
deps = Deps(client=client, brave_api_key=brave_api_key)
result = await agent.run(
user_input, deps=deps
)
print('Response:', result.data)
if __name__ == '__main__':
asyncio.run(main()) |