mirror of
https://github.com/rowboatlabs/rowboat.git
synced 2026-04-26 08:56:22 +02:00
add web search tool in run_streamed
This commit is contained in:
arkml
Ramnique Singh
parent
6312fe8e83
commit
8e16ac6204
1 changed files with 228 additions and 25 deletions
|
|
@ -282,6 +282,54 @@ async def run_turn_streamed(
|
|||
print('-'*50)
|
||||
print(f"Found usage information. Updated cumulative tokens: {tokens_used}")
|
||||
print('-'*50)
|
||||
|
||||
# Handle ResponseFunctionWebSearch specifically
|
||||
if hasattr(event, 'data') and hasattr(event.data, 'raw_item'):
|
||||
raw_item = event.data.raw_item
|
||||
|
||||
# Check if it's a web search call
|
||||
if (hasattr(raw_item, 'type') and raw_item.type == 'web_search_call') or (
|
||||
isinstance(raw_item, dict) and raw_item.get('type') == 'web_search_call'
|
||||
):
|
||||
# Get call_id safely, regardless of structure
|
||||
call_id = None
|
||||
if hasattr(raw_item, 'id'):
|
||||
call_id = raw_item.id
|
||||
elif isinstance(raw_item, dict) and 'id' in raw_item:
|
||||
call_id = raw_item['id']
|
||||
else:
|
||||
call_id = str(uuid.uuid4())
|
||||
|
||||
# Get status safely
|
||||
status = 'unknown'
|
||||
if hasattr(raw_item, 'status'):
|
||||
status = raw_item.status
|
||||
elif isinstance(raw_item, dict) and 'status' in raw_item:
|
||||
status = raw_item['status']
|
||||
|
||||
# Emit a tool call for web search
|
||||
message = {
|
||||
'content': None,
|
||||
'role': 'assistant',
|
||||
'sender': current_agent.name if current_agent else None,
|
||||
'tool_calls': [{
|
||||
'function': {
|
||||
'name': 'web_search',
|
||||
'arguments': json.dumps({
|
||||
'search_id': call_id,
|
||||
'status': status
|
||||
})
|
||||
},
|
||||
'id': call_id,
|
||||
'type': 'function'
|
||||
}],
|
||||
'tool_call_id': None,
|
||||
'tool_name': None,
|
||||
'response_type': 'internal'
|
||||
}
|
||||
print("Yielding web search raw response message: ", message)
|
||||
yield ('message', message)
|
||||
|
||||
continue
|
||||
|
||||
# Update current agent when it changes
|
||||
|
|
@ -334,45 +382,112 @@ async def run_turn_streamed(
|
|||
elif event.type == "run_item_stream_event":
|
||||
current_agent = event.item.agent
|
||||
if event.item.type == "tool_call_item":
|
||||
message = {
|
||||
'content': None,
|
||||
'role': 'assistant',
|
||||
'sender': current_agent.name if current_agent else None,
|
||||
'tool_calls': [{
|
||||
'function': {
|
||||
'name': event.item.raw_item.name,
|
||||
'arguments': event.item.raw_item.arguments
|
||||
},
|
||||
'id': event.item.raw_item.call_id,
|
||||
'type': 'function'
|
||||
}],
|
||||
'tool_call_id': None,
|
||||
'tool_name': None,
|
||||
'response_type': 'internal'
|
||||
}
|
||||
# Check if it's a ResponseFunctionWebSearch object
|
||||
if hasattr(event.item.raw_item, 'type') and event.item.raw_item.type == 'web_search_call':
|
||||
call_id = event.item.raw_item.id if hasattr(event.item.raw_item, 'id') else str(uuid.uuid4())
|
||||
message = {
|
||||
'content': None,
|
||||
'role': 'assistant',
|
||||
'sender': current_agent.name if current_agent else None,
|
||||
'tool_calls': [{
|
||||
'function': {
|
||||
'name': 'web_search',
|
||||
'arguments': json.dumps({
|
||||
'search_id': call_id
|
||||
})
|
||||
},
|
||||
'id': call_id,
|
||||
'type': 'function'
|
||||
}],
|
||||
'tool_call_id': None,
|
||||
'tool_name': None,
|
||||
'response_type': 'internal'
|
||||
}
|
||||
else:
|
||||
# Handle normal tool calls
|
||||
message = {
|
||||
'content': None,
|
||||
'role': 'assistant',
|
||||
'sender': current_agent.name if current_agent else None,
|
||||
'tool_calls': [{
|
||||
'function': {
|
||||
'name': event.item.raw_item.name,
|
||||
'arguments': event.item.raw_item.arguments
|
||||
},
|
||||
'id': event.item.raw_item.call_id,
|
||||
'type': 'function'
|
||||
}],
|
||||
'tool_call_id': None,
|
||||
'tool_name': None,
|
||||
'response_type': 'internal'
|
||||
}
|
||||
print("Yielding message: ", message)
|
||||
yield ('message', message)
|
||||
|
||||
elif event.item.type == "tool_call_output_item":
|
||||
message = {
|
||||
'content': str(event.item.output),
|
||||
'role': 'tool',
|
||||
'sender': None,
|
||||
'tool_calls': None,
|
||||
'tool_call_id': event.item.raw_item['call_id'],
|
||||
'tool_name': event.item.raw_item.get('name', None),
|
||||
'response_type': 'internal'
|
||||
}
|
||||
# Check if it's a web search result
|
||||
if isinstance(event.item.raw_item, dict) and event.item.raw_item.get('type') == 'web_search_results':
|
||||
call_id = event.item.raw_item.get('search_id', event.item.raw_item.get('id', str(uuid.uuid4())))
|
||||
message = {
|
||||
'content': str(event.item.output),
|
||||
'role': 'tool',
|
||||
'sender': None,
|
||||
'tool_calls': None,
|
||||
'tool_call_id': call_id,
|
||||
'tool_name': 'web_search',
|
||||
'response_type': 'internal'
|
||||
}
|
||||
else:
|
||||
# Safe extraction of call_id and name
|
||||
call_id = None
|
||||
tool_name = None
|
||||
|
||||
# Handle different types of raw_item
|
||||
if isinstance(event.item.raw_item, dict):
|
||||
call_id = event.item.raw_item.get('call_id')
|
||||
tool_name = event.item.raw_item.get('name')
|
||||
elif hasattr(event.item.raw_item, 'call_id'):
|
||||
call_id = event.item.raw_item.call_id
|
||||
if hasattr(event.item.raw_item, 'name'):
|
||||
tool_name = event.item.raw_item.name
|
||||
|
||||
message = {
|
||||
'content': str(event.item.output),
|
||||
'role': 'tool',
|
||||
'sender': None,
|
||||
'tool_calls': None,
|
||||
'tool_call_id': call_id,
|
||||
'tool_name': tool_name,
|
||||
'response_type': 'internal'
|
||||
}
|
||||
|
||||
print("Yielding message: ", message)
|
||||
yield ('message', message)
|
||||
|
||||
elif event.item.type == "message_output_item":
|
||||
content = ""
|
||||
url_citations = []
|
||||
|
||||
# Extract text content and any URL citations
|
||||
if hasattr(event.item.raw_item, 'content'):
|
||||
for content_item in event.item.raw_item.content:
|
||||
# Handle text content
|
||||
if hasattr(content_item, 'text'):
|
||||
content += content_item.text
|
||||
|
||||
# Extract URL citations if present
|
||||
if hasattr(content_item, 'annotations'):
|
||||
for annotation in content_item.annotations:
|
||||
if hasattr(annotation, 'type') and annotation.type == 'url_citation':
|
||||
citation = {
|
||||
'url': annotation.url if hasattr(annotation, 'url') else '',
|
||||
'title': annotation.title if hasattr(annotation, 'title') else '',
|
||||
'start_index': annotation.start_index if hasattr(annotation, 'start_index') else 0,
|
||||
'end_index': annotation.end_index if hasattr(annotation, 'end_index') else 0
|
||||
}
|
||||
url_citations.append(citation)
|
||||
|
||||
# Create message with URL citations if they exist
|
||||
message = {
|
||||
'content': content,
|
||||
'role': 'assistant',
|
||||
|
|
@ -382,9 +497,97 @@ async def run_turn_streamed(
|
|||
'tool_name': None,
|
||||
'response_type': 'external'
|
||||
}
|
||||
|
||||
# Add citations if any were found
|
||||
if url_citations:
|
||||
message['citations'] = url_citations
|
||||
|
||||
print("Yielding message: ", message)
|
||||
yield ('message', message)
|
||||
|
||||
# Handle web search function call events
|
||||
elif event.item.type == "web_search_call_item" or (hasattr(event.item, 'raw_item') and hasattr(event.item.raw_item, 'type') and event.item.raw_item.type == 'web_search_call'):
|
||||
# Extract web search call ID if available
|
||||
call_id = None
|
||||
if hasattr(event.item.raw_item, 'id'):
|
||||
call_id = event.item.raw_item.id
|
||||
|
||||
message = {
|
||||
'content': None,
|
||||
'role': 'assistant',
|
||||
'sender': current_agent.name if current_agent else None,
|
||||
'tool_calls': [{
|
||||
'function': {
|
||||
'name': 'web_search',
|
||||
'arguments': json.dumps({
|
||||
'search_id': call_id
|
||||
})
|
||||
},
|
||||
'id': call_id or str(uuid.uuid4()),
|
||||
'type': 'function'
|
||||
}],
|
||||
'tool_call_id': None,
|
||||
'tool_name': None,
|
||||
'response_type': 'internal'
|
||||
}
|
||||
print("Yielding web search message: ", message)
|
||||
yield ('message', message)
|
||||
|
||||
# Handle web search results
|
||||
elif event.item.type == "web_search_results_item" or (
|
||||
hasattr(event.item, 'raw_item') and (
|
||||
(hasattr(event.item.raw_item, 'type') and event.item.raw_item.type == 'web_search_results') or
|
||||
(isinstance(event.item.raw_item, dict) and event.item.raw_item.get('type') == 'web_search_results')
|
||||
)
|
||||
):
|
||||
# Extract call_id safely
|
||||
call_id = None
|
||||
raw_item = event.item.raw_item
|
||||
|
||||
# Try several ways to get the search_id or id
|
||||
if hasattr(raw_item, 'search_id'):
|
||||
call_id = raw_item.search_id
|
||||
elif isinstance(raw_item, dict) and 'search_id' in raw_item:
|
||||
call_id = raw_item['search_id']
|
||||
elif hasattr(raw_item, 'id'):
|
||||
call_id = raw_item.id
|
||||
elif isinstance(raw_item, dict) and 'id' in raw_item:
|
||||
call_id = raw_item['id']
|
||||
else:
|
||||
call_id = str(uuid.uuid4())
|
||||
|
||||
# Extract results content safely
|
||||
results = {}
|
||||
|
||||
# Try event.item.output first
|
||||
if hasattr(event.item, 'output'):
|
||||
results = event.item.output
|
||||
# Then try raw_item.results
|
||||
elif hasattr(raw_item, 'results'):
|
||||
results = raw_item.results
|
||||
elif isinstance(raw_item, dict) and 'results' in raw_item:
|
||||
results = raw_item['results']
|
||||
|
||||
# Format the results for output
|
||||
results_str = ""
|
||||
try:
|
||||
results_str = json.dumps(results) if results else ""
|
||||
except Exception as e:
|
||||
print(f"Error serializing results: {str(e)}")
|
||||
results_str = str(results)
|
||||
|
||||
message = {
|
||||
'content': results_str,
|
||||
'role': 'tool',
|
||||
'sender': None,
|
||||
'tool_calls': None,
|
||||
'tool_call_id': call_id,
|
||||
'tool_name': 'web_search',
|
||||
'response_type': 'internal'
|
||||
}
|
||||
print("Yielding web search results: ", message)
|
||||
yield ('message', message)
|
||||
|
||||
print(f"\n{'='*50}\n")
|
||||
|
||||
# After all events are processed, set final state
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue