exa-py 1.0.9__py3-none-any.whl → 1.0.11__py3-none-any.whl

exa_py/api.py

@@ -1,9 +1,12 @@
 from __future__ import annotations
 from dataclasses import dataclass
 import dataclasses
+from functools import wraps
 import re
 import requests
 from typing import (
+    Callable,
+    Iterable,
     List,
     Optional,
     Dict,
@@ -15,6 +18,28 @@ from typing import (
 )
 from typing_extensions import TypedDict
 
+import httpx
+from openai import NOT_GIVEN, NotGiven, OpenAI
+from openai.types.chat.chat_completion_stream_options_param import (
+    ChatCompletionStreamOptionsParam,
+)
+from openai.types.chat.chat_completion_message_param import ChatCompletionMessageParam
+from openai.types.chat_model import ChatModel
+from openai.types.chat.chat_completion_tool_choice_option_param import (
+    ChatCompletionToolChoiceOptionParam,
+)
+from openai.types.chat.chat_completion_tool_param import ChatCompletionToolParam
+from openai._types import Headers, Query, Body
+from openai.types.chat import completion_create_params
+from exa_py.utils import (
+    ExaOpenAICompletion,
+    add_message_to_messages,
+    format_exa_result,
+    maybe_get_query,
+)
+
+
+
 
 def snake_to_camel(snake_str: str) -> str:
     """Convert snake_case string to camelCase.
@@ -319,7 +344,7 @@ class Exa:
         self,
         api_key: Optional[str],
         base_url: str = "https://api.exa.ai",
-        user_agent: str = "exa-py 1.0.9",
+        user_agent: str = "exa-py 1.0.11",
     ):
         """Initialize the Exa client with the provided API key and optional base URL and user agent.
 
@@ -646,3 +671,177 @@ class Exa:
             [Result(**to_snake_case(result)) for result in data["results"]],
             data["autopromptString"] if "autopromptString" in data else None,
         )
+    def wrap(self, client: OpenAI):
+        """Wrap an OpenAI client with Exa functionality.
+
+        After wrapping, any call to `client.chat.completions.create` will be intercepted and enhanced with Exa functionality.
+
+        To disable Exa functionality for a specific call, set `use_exa="none"` in the call to `client.chat.completions.create`.
+
+        Args:
+            client (OpenAI): The OpenAI client to wrap.
+
+        Returns:
+            OpenAI: The wrapped OpenAI client.
+        """
+
+        func = client.chat.completions.create
+
+        @wraps(func)
+        def create_with_rag(
+            # Mandatory OpenAI args
+            messages: Iterable[ChatCompletionMessageParam],
+            model: Union[str, ChatModel],
+            # Optional OpenAI args
+            frequency_penalty: Optional[float] | NotGiven = NOT_GIVEN,
+            function_call: completion_create_params.FunctionCall | NotGiven = NOT_GIVEN,
+            functions: (
+                Iterable[completion_create_params.Function] | NotGiven
+            ) = NOT_GIVEN,
+            logit_bias: Optional[Dict[str, int]] | NotGiven = NOT_GIVEN,
+            logprobs: Optional[bool] | NotGiven = NOT_GIVEN,
+            max_tokens: Optional[int] | NotGiven = NOT_GIVEN,
+            n: Optional[int] | NotGiven = NOT_GIVEN,
+            presence_penalty: Optional[float] | NotGiven = NOT_GIVEN,
+            response_format: (
+                completion_create_params.ResponseFormat | NotGiven
+            ) = NOT_GIVEN,
+            seed: Optional[int] | NotGiven = NOT_GIVEN,
+            stop: Union[Optional[str], List[str]] | NotGiven = NOT_GIVEN,
+            stream: Optional[Literal[False]] | NotGiven = NOT_GIVEN,
+            stream_options: (
+                Optional[ChatCompletionStreamOptionsParam] | NotGiven
+            ) = NOT_GIVEN,
+            temperature: Optional[float] | NotGiven = NOT_GIVEN,
+            tool_choice: ChatCompletionToolChoiceOptionParam | NotGiven = NOT_GIVEN,
+            tools: Iterable[ChatCompletionToolParam] | NotGiven = NOT_GIVEN,
+            top_logprobs: Optional[int] | NotGiven = NOT_GIVEN,
+            top_p: Optional[float] | NotGiven = NOT_GIVEN,
+            user: str | NotGiven = NOT_GIVEN,
+            # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
+            # The extra values given here take precedence over values defined on the client or passed to this method.
+            extra_headers: Headers | None = None,
+            extra_query: Query | None = None,
+            extra_body: Body | None = None,
+            timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN,
+            # Exa args
+            use_exa: Optional[Literal["required", "none", "auto"]] = "auto",
+            highlights: Union[HighlightsContentsOptions, Literal[True], None] = None,
+            num_results: Optional[int] = 3,
+            include_domains: Optional[List[str]] = None,
+            exclude_domains: Optional[List[str]] = None,
+            start_crawl_date: Optional[str] = None,
+            end_crawl_date: Optional[str] = None,
+            start_published_date: Optional[str] = None,
+            end_published_date: Optional[str] = None,
+            use_autoprompt: Optional[bool] = True,
+            type: Optional[str] = None,
+            category: Optional[str] = None,
+            result_max_len: int = 2048,
+        ):
+            exa_kwargs = {
+                "num_results": num_results,
+                "include_domains": include_domains,
+                "exclude_domains": exclude_domains,
+                "highlights": highlights,
+                "start_crawl_date": start_crawl_date,
+                "end_crawl_date": end_crawl_date,
+                "start_published_date": start_published_date,
+                "end_published_date": end_published_date,
+                "use_autoprompt": use_autoprompt,
+                "type": type,
+                "category": category,
+            }
+
+            create_kwargs = {
+                "model": model,
+                "frequency_penalty": frequency_penalty,
+                "function_call": function_call,
+                "functions": functions,
+                "logit_bias": logit_bias,
+                "logprobs": logprobs,
+                "max_tokens": max_tokens,
+                "n": n,
+                "presence_penalty": presence_penalty,
+                "response_format": response_format,
+                "seed": seed,
+                "stop": stop,
+                "stream": stream,
+                "stream_options": stream_options,
+                "temperature": temperature,
+                "tool_choice": tool_choice,
+                "tools": tools,
+                "top_logprobs": top_logprobs,
+                "top_p": top_p,
+                "user": user,
+                "extra_headers": extra_headers,
+                "extra_query": extra_query,
+                "extra_body": extra_body,
+                "timeout": timeout,
+            }
+
+            if use_exa != "none":
+                assert tools is NOT_GIVEN, "Tool use is not supported with Exa"
+                create_kwargs["tool_choice"] = use_exa
+
+            return self._create_with_tool(
+                create_fn=func,
+                messages=list(messages),
+                max_len=result_max_len,
+                create_kwargs=create_kwargs,
+                exa_kwargs=exa_kwargs,
+            )
+
+        print("Wrapping OpenAI client with Exa functionality.", type(create_with_rag))
+        client.chat.completions.create = create_with_rag # type: ignore
+
+        return client
+
+    def _create_with_tool(
+        self,
+        create_fn: Callable,
+        messages: List[ChatCompletionMessageParam],
+        max_len,
+        create_kwargs,
+        exa_kwargs,
+    ) -> ExaOpenAICompletion:
+        tools = [
+            {
+                "type": "function",
+                "function": {
+                    "name": "search",
+                    "description": "Search the web for relevant information.",
+                    "parameters": {
+                        "type": "object",
+                        "properties": {
+                            "query": {
+                                "type": "string",
+                                "description": "The query to search for.",
+                            },
+                        },
+                        "required": ["query"],
+                    },
+                },
+            }
+        ]
+
+        create_kwargs["tools"] = tools
+
+        completion = create_fn(messages=messages, **create_kwargs)
+    
+        query = maybe_get_query(completion)
+
+        if not query:
+            return ExaOpenAICompletion.from_completion(completion=completion, exa_result=None)
+
+        exa_result = self.search_and_contents(query, **exa_kwargs)
+        exa_str = format_exa_result(exa_result, max_len=max_len)
+        new_messages = add_message_to_messages(completion, messages, exa_str)
+        # For now, don't allow recursive tool calls
+        create_kwargs["tool_choice"] = "none"
+        completion = create_fn(messages=new_messages, **create_kwargs)
+
+        exa_completion = ExaOpenAICompletion.from_completion(
+            completion=completion, exa_result=exa_result
+        )
+        return exa_completion

exa_py/utils.py

@@ -0,0 +1,78 @@
+import json
+from typing import Optional
+from openai.types.chat import ChatCompletion
+
+from typing import TYPE_CHECKING 
+if TYPE_CHECKING:
+    from exa_py.api import ResultWithText, SearchResponse
+
+
+
+def maybe_get_query(completion) -> str | None:
+    """Extract query from completion if it exists."""
+    if completion.choices[0].message.tool_calls:
+        for tool_call in completion.choices[0].message.tool_calls:
+            if tool_call.function.name == "search":
+                query = json.loads(tool_call.function.arguments)["query"]
+                return query
+    return None
+
+
+def add_message_to_messages(completion, messages, exa_result) -> list[dict]:
+    """Add assistant message and exa result to messages list. Also remove previous exa call and results."""
+    assistant_message = completion.choices[0].message
+    assert assistant_message.tool_calls, "Must use this with a tool call request"
+    # Remove previous exa call and results to prevent blowing up history
+    messages = [
+        message
+        for message in messages
+        if not (message.get("role") == "function")
+    ]
+    
+    messages.extend([
+        assistant_message,
+        {
+            "role": "tool",
+            "name": "search",
+            "tool_call_id": assistant_message.tool_calls[0].id,
+            "content": exa_result,
+        }
+    ])
+
+    return messages
+
+
+def format_exa_result(exa_result, max_len: int=-1):
+    """Format exa result for pasting into chat."""
+    str = [
+        f"Url: {result.url}\nTitle: {result.title}\n{result.text[:max_len]}\n"
+        for result in exa_result.results
+    ]
+
+    return "\n".join(str)
+
+
+class ExaOpenAICompletion(ChatCompletion):
+    """Exa wrapper for OpenAI completion."""
+    def __init__(self, exa_result: Optional["SearchResponse[ResultWithText]"], **kwargs):
+        super().__init__(**kwargs)
+        self.exa_result = exa_result
+    
+
+    @classmethod
+    def from_completion(
+        cls, 
+        exa_result: Optional["SearchResponse[ResultWithText]"], 
+        completion: ChatCompletion
+    ):
+
+        return cls(
+            exa_result=exa_result,
+            id=completion.id,
+            choices=completion.choices,
+            created=completion.created,
+            model=completion.model,
+            object=completion.object,
+            system_fingerprint=completion.system_fingerprint,
+            usage=completion.usage,
+        )

{exa_py-1.0.9.dist-info → exa_py-1.0.11.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: exa-py
-Version: 1.0.9
+Version: 1.0.11
 Summary: Python SDK for Exa API.
 Home-page: https://github.com/exa-labs/exa-py
 Author: Exa
@@ -18,6 +18,9 @@ Classifier: Programming Language :: Python :: 3.12
 Description-Content-Type: text/markdown
 Requires-Dist: requests
 Requires-Dist: typing-extensions
+Requires-Dist: openai
+Provides-Extra: openai
+Requires-Dist: openai ; extra == 'openai'
 
 # Exa
 

exa_py-1.0.11.dist-info/RECORD

@@ -0,0 +1,8 @@
+exa_py/__init__.py,sha256=aVF1zB_UV3dagJ5Vn2WrdcInzibdIW61M89sjwRCU_g,29
+exa_py/api.py,sha256=l1gg_Sp2c9R9SpUpPI7G6M8mrnYULfsxZBv8lHVoxPU,31032
+exa_py/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+exa_py/utils.py,sha256=pF91TKGGg4SbSVrdWroQze90-D9AlvqLoKWlHIeyLaE,2402
+exa_py-1.0.11.dist-info/METADATA,sha256=jdepNNb0qBPgzPzHZCQ33Ocyfe0dIe4A3DhppxbwQQQ,3137
+exa_py-1.0.11.dist-info/WHEEL,sha256=pkctZYzUS4AYVn6dJ-7367OJZivF2e8RA9b_ZBjif18,92
+exa_py-1.0.11.dist-info/top_level.txt,sha256=Mfkmscdw9HWR1PtVhU1gAiVo6DHu_tyiVdb89gfZBVI,7
+exa_py-1.0.11.dist-info/RECORD,,

exa_py-1.0.9.dist-info/RECORD

@@ -1,7 +0,0 @@
-exa_py/__init__.py,sha256=aVF1zB_UV3dagJ5Vn2WrdcInzibdIW61M89sjwRCU_g,29
-exa_py/api.py,sha256=wV3JOtEDN6yggi4TS2XaAhcvHU85LIRine_mioZ049E,22932
-exa_py/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-exa_py-1.0.9.dist-info/METADATA,sha256=XqcyExdupc9LOfn595tOlb3Y7LLHdhY7yQ08j4UPnz0,3049
-exa_py-1.0.9.dist-info/WHEEL,sha256=pkctZYzUS4AYVn6dJ-7367OJZivF2e8RA9b_ZBjif18,92
-exa_py-1.0.9.dist-info/top_level.txt,sha256=Mfkmscdw9HWR1PtVhU1gAiVo6DHu_tyiVdb89gfZBVI,7
-exa_py-1.0.9.dist-info/RECORD,,