OpenAI Responses API

import json
from openai import OpenAI
from exa_py import Exa

OPENAI_API_KEY = ""  # Add your OpenAI API key here
EXA_API_KEY = ""     # Add your Exa API key here

# Define the tool for Exa web search
tools = [{
    "type": "function",
    "name": "exa_websearch",
    "description": "Search the web using Exa. Provide relevant links in your answer.",
    "parameters": {
        "type": "object",
        "properties": {
            "query": {
                "type": "string",
                "description": "Search query for Exa."
            }
        },
        "required": ["query"],
        "additionalProperties": False
    },
    "strict": True
}]

# Define the system message
system_message = {"role": "system", "content": "You are a helpful assistant. Use exa_websearch to find info when relevant. Always list sources."}

def run_exa_search(user_query):
    """Run an Exa web search with a dynamic user query."""
    openai_client = OpenAI(api_key=OPENAI_API_KEY)
    exa = Exa(api_key=EXA_API_KEY)
    
    # Create messages with the dynamic user query
    messages = [
        system_message,
        {"role": "user", "content": user_query}
    ]

    # Send initial request
    print("Sending initial request to OpenAI...")
    response = openai_client.responses.create(
        model="gpt-4o",
        input=messages,
        tools=tools
    )
    print("Initial OpenAI response:", response.output)

    # Check if the model returned a function call
    function_call = None
    for item in response.output:
        if item.type == "function_call" and item.name == "exa_websearch":
            function_call = item
            break

    # If exa_websearch was called
    if function_call:
        call_id = function_call.call_id
        args = json.loads(function_call.arguments)
        query = args.get("query", "")

        print(f"\nOpenAI requested a web search for: {query}")
        search_results = exa.search_and_contents(
            query=query,
            text = {
              "max_characters": 4000
            },
            type="auto"
        )
        
        # Store citations for later use in formatting
        citations = [{"url": result.url, "title": result.title} for result in search_results.results]

        search_results_str = str(search_results)

        # Provide the function call + function_call_output to the conversation
        messages.append({
            "type": "function_call",
            "name": function_call.name,
            "arguments": function_call.arguments,
            "call_id": call_id
        })
        messages.append({
            "type": "function_call_output",
            "call_id": call_id,
            "output": search_results_str
        })

        print("\nSending search results back to OpenAI for a final answer...")
        response = openai_client.responses.create(
            model="gpt-4o",
            input=messages,
            tools=tools
        )
        
        # Format the final response to include citations
        if hasattr(response, 'output_text') and response.output_text:
            # Add citations to the final output
            formatted_response = format_response_with_citations(response.output_text, citations)
            
            # Create a custom response object with citations
            if hasattr(response, 'model_dump'):
                # For newer versions of the OpenAI library that use Pydantic
                response_dict = response.model_dump()
            else:
                # For older versions or if model_dump is not available
                response_dict = response.dict() if hasattr(response, 'dict') else response.__dict__
            
            # Update the output with annotations
            if response.output and len(response.output) > 0:
                response_dict['output'] = [{
                    "type": "message",
                    "id": response.output[0].id if hasattr(response.output[0], 'id') else "msg_custom",
                    "status": "completed",
                    "role": "assistant",
                    "content": [{
                        "type": "output_text",
                        "text": formatted_response["text"],
                        "annotations": formatted_response["annotations"]
                    }]
                }]
                
                # Update the output_text property
                response_dict['output_text'] = formatted_response["text"]
                
                # Create a new response object (implementation may vary based on the OpenAI SDK version)
                try:
                    response = type(response)(**response_dict)
                except:
                    # If we can't create a new instance, we'll just print the difference
                    print("\nFormatted response with citations would be:", formatted_response)

    # Print final answer text
    print("\nFinal Answer:\n", response.output_text)
    print("\nAnnotations:", json.dumps(response.output[0].content[0].annotations if hasattr(response, 'output') and response.output and hasattr(response.output[0], 'content') else [], indent=2))
    print("\nFull Response with Citations:", response)
    
    return response

def format_response_with_citations(text, citations):
    """Format the response to include citations as annotations."""
    annotations = []
    formatted_text = text
    
    # For each citation, append a numbered reference to the text
    for i, citation in enumerate(citations):
        # Create annotation object
        start_index = len(formatted_text)
        citation_text = f"\n\n[{i+1}] {citation['url']}"
        end_index = start_index + len(citation_text)
        
        annotation = {
            "type": "url_citation",
            "start_index": start_index,
            "end_index": end_index,
            "url": citation["url"],
            "title": citation["title"]
        }
        
        # Add annotation to the array
        annotations.append(annotation)
        
        # Append citation to text
        formatted_text += citation_text
    
    return {
        "text": formatted_text,
        "annotations": annotations
    }


if __name__ == "__main__":
    # Example of how to use with a dynamic query
    user_query = input("Enter your question: ")
    run_exa_search(user_query)