diff --git a/README.md b/README.md index 7d2951f..33211f3 100644 --- a/README.md +++ b/README.md @@ -293,6 +293,29 @@ response = await client.search("latest AI research") - Custom session proxies take precedence over SDK proxy settings - The SDK will **not** close externally-provided sessions — you manage the lifecycle +## Session & User Tracking + +`session_id`, `human_id`, and `client_name` are optional identifiers that help attribute requests to a logical session, an end user, and a named client. All three are sent as HTTP headers (`X-Session-Id`, `X-Human-Id`, `X-Client-Name`) and are never persisted in raw form — `human_id` is hashed server-side. + +Set them once at client init, or per-call (per-call wins): + +```python +from tavily import TavilyClient + +# Client-level — applied to every request +client = TavilyClient( + api_key="tvly-YOUR_API_KEY", + session_id="my-session-123", + human_id="internal-user-id-42", + client_name="my-app", +) + +# Per-call override +client.search("hello", session_id="ad-hoc-session") +``` + +All three are opt-in. Leave them unset and the SDK sends nothing — behavior is identical to earlier versions. + ## Documentation For a complete guide on how to use the different endpoints and their parameters, please head to our [Python API Reference](https://docs.tavily.com/sdk/python/reference). diff --git a/setup.py b/setup.py index 4c4618f..3023eec 100644 --- a/setup.py +++ b/setup.py @@ -5,7 +5,7 @@ setup( name='tavily-python', - version='0.7.23', + version='0.7.24', url='https://github.com/tavily-ai/tavily-python', author='Tavily AI', author_email='support@tavily.com', diff --git a/tavily/async_tavily.py b/tavily/async_tavily.py index dc896b8..663d228 100644 --- a/tavily/async_tavily.py +++ b/tavily/async_tavily.py @@ -20,6 +20,9 @@ def __init__(self, api_key: Optional[str] = None, api_base_url: Optional[str] = None, client_source: Optional[str] = None, project_id: Optional[str] = None, + session_id: Optional[str] = None, + human_id: Optional[str] = None, + client_name: Optional[str] = None, client: Optional[httpx.AsyncClient] = None): if api_key is None: api_key = os.getenv("TAVILY_API_KEY") @@ -36,7 +39,10 @@ def __init__(self, api_key: Optional[str] = None, "Content-Type": "application/json", **({"Authorization": f"Bearer {api_key}"} if api_key else {}), "X-Client-Source": client_source or "tavily-python", - **({"X-Project-ID": tavily_project} if tavily_project else {}) + **({"X-Project-ID": tavily_project} if tavily_project else {}), + **({"X-Session-Id": session_id} if session_id else {}), + **({"X-Human-Id": human_id} if human_id else {}), + **({"X-Client-Name": client_name} if client_name else {}), } self._external_client = client is not None @@ -83,6 +89,23 @@ async def __aenter__(self): async def __aexit__(self, exc_type, exc_val, exc_tb): await self.close() + @staticmethod + def _pop_request_headers(kwargs: dict) -> Optional[dict]: + """Pop session_id, human_id, and client_name from kwargs and return them as headers. + + Returns None when no overrides are provided so callers can omit the headers kwarg. + """ + overrides = {} + for key, header_name in ( + ("session_id", "X-Session-Id"), + ("human_id", "X-Human-Id"), + ("client_name", "X-Client-Name"), + ): + value = kwargs.pop(key, None) + if value is not None: + overrides[header_name] = str(value) + return overrides or None + async def _search( self, query: str, @@ -132,13 +155,14 @@ async def _search( data = {k: v for k, v in data.items() if v is not None} + override_headers = self._pop_request_headers(kwargs) if kwargs: data.update(kwargs) timeout = min(timeout, 120) try: - response = await self._client.post("/search", content=json.dumps(data), timeout=timeout) + response = await self._client.post("/search", content=json.dumps(data), timeout=timeout, **({"headers": override_headers} if override_headers else {})) except httpx.TimeoutException: raise TimeoutError(timeout) @@ -247,11 +271,12 @@ async def _extract( data = {k: v for k, v in data.items() if v is not None} + override_headers = self._pop_request_headers(kwargs) if kwargs: data.update(kwargs) try: - response = await self._client.post("/extract", content=json.dumps(data), timeout=timeout) + response = await self._client.post("/extract", content=json.dumps(data), timeout=timeout, **({"headers": override_headers} if override_headers else {})) except httpx.TimeoutException: raise TimeoutError(timeout) @@ -355,13 +380,14 @@ async def _crawl(self, "chunks_per_source": chunks_per_source, } + override_headers = self._pop_request_headers(kwargs) if kwargs: data.update(kwargs) data = {k: v for k, v in data.items() if v is not None} try: - response = await self._client.post("/crawl", content=json.dumps(data), timeout=timeout) + response = await self._client.post("/crawl", content=json.dumps(data), timeout=timeout, **({"headers": override_headers} if override_headers else {})) except httpx.TimeoutException: raise TimeoutError(timeout) @@ -465,13 +491,14 @@ async def _map(self, "include_usage": include_usage, } + override_headers = self._pop_request_headers(kwargs) if kwargs: data.update(kwargs) data = {k: v for k, v in data.items() if v is not None} try: - response = await self._client.post("/map", content=json.dumps(data), timeout=timeout) + response = await self._client.post("/map", content=json.dumps(data), timeout=timeout, **({"headers": override_headers} if override_headers else {})) except httpx.TimeoutException: raise TimeoutError(timeout) @@ -659,6 +686,7 @@ def _research(self, data = {k: v for k, v in data.items() if v is not None} + override_headers = self._pop_request_headers(kwargs) if kwargs: data.update(kwargs) @@ -669,7 +697,8 @@ async def stream_generator() -> AsyncGenerator[bytes, None]: "POST", "/research", content=json.dumps(data), - timeout=timeout + timeout=timeout, + **({"headers": override_headers} if override_headers else {}) ) as response: if response.status_code != 200: try: @@ -701,7 +730,7 @@ async def stream_generator() -> AsyncGenerator[bytes, None]: else: async def _make_request(): try: - response = await self._client.post("/research", content=json.dumps(data), timeout=timeout) + response = await self._client.post("/research", content=json.dumps(data), timeout=timeout, **({"headers": override_headers} if override_headers else {})) except httpx.TimeoutException: raise TimeoutError(timeout) diff --git a/tavily/tavily.py b/tavily/tavily.py index b059734..86ba02d 100644 --- a/tavily/tavily.py +++ b/tavily/tavily.py @@ -11,7 +11,18 @@ class TavilyClient: Tavily API client class. """ - def __init__(self, api_key: Optional[str] = None, proxies: Optional[dict[str, str]] = None, api_base_url: Optional[str] = None, client_source: Optional[str] = None, project_id: Optional[str] = None, session: Optional[requests.Session] = None): + def __init__( + self, + api_key: Optional[str] = None, + proxies: Optional[dict[str, str]] = None, + api_base_url: Optional[str] = None, + client_source: Optional[str] = None, + project_id: Optional[str] = None, + session_id: Optional[str] = None, + human_id: Optional[str] = None, + client_name: Optional[str] = None, + session: Optional[requests.Session] = None, + ): if api_key is None: api_key = os.getenv("TAVILY_API_KEY") @@ -25,16 +36,19 @@ def __init__(self, api_key: Optional[str] = None, proxies: Optional[dict[str, st resolved_proxies = {k: v for k, v in resolved_proxies.items() if v} or None tavily_project = project_id or os.getenv("TAVILY_PROJECT") - + self.base_url = api_base_url or "https://api.tavily.com" self.api_key = api_key self.proxies = resolved_proxies - + self.headers = { "Content-Type": "application/json", **({"Authorization": f"Bearer {self.api_key}"} if self.api_key else {}), "X-Client-Source": client_source or "tavily-python", - **({"X-Project-ID": tavily_project} if tavily_project else {}) + **({"X-Project-ID": tavily_project} if tavily_project else {}), + **({"X-Session-Id": session_id} if session_id else {}), + **({"X-Human-Id": human_id} if human_id else {}), + **({"X-Client-Name": client_name} if client_name else {}), } self._external_session = session is not None @@ -59,6 +73,23 @@ def __enter__(self): def __exit__(self, exc_type, exc_val, exc_tb): self.close() + @staticmethod + def _pop_request_headers(kwargs: dict) -> Optional[dict]: + """Pop session_id, human_id, and client_name from kwargs and return them as headers. + + Returns None when no overrides are provided so callers can omit the headers kwarg. + """ + overrides = {} + for key, header_name in ( + ("session_id", "X-Session-Id"), + ("human_id", "X-Human-Id"), + ("client_name", "X-Client-Name"), + ): + value = kwargs.pop(key, None) + if value is not None: + overrides[header_name] = str(value) + return overrides or None + def _search(self, query: str, search_depth: Literal["basic", "advanced", "fast", "ultra-fast"] = None, @@ -108,6 +139,7 @@ def _search(self, data = {k: v for k, v in data.items() if v is not None} + override_headers = self._pop_request_headers(kwargs) if kwargs: data.update(kwargs) @@ -116,7 +148,7 @@ def _search(self, payload = json.dumps(data) try: - response = self.session.post(url, data=payload, timeout=timeout) + response = self.session.post(url, data=payload, timeout=timeout, **({"headers": override_headers} if override_headers else {})) except requests.exceptions.Timeout: raise TimeoutError(timeout) @@ -219,11 +251,12 @@ def _extract(self, data = {k: v for k, v in data.items() if v is not None} + override_headers = self._pop_request_headers(kwargs) if kwargs: data.update(kwargs) try: - response = self.session.post(self.base_url + "/extract", data=json.dumps(data), timeout=timeout) + response = self.session.post(self.base_url + "/extract", data=json.dumps(data), timeout=timeout, **({"headers": override_headers} if override_headers else {})) except requests.exceptions.Timeout: raise TimeoutError(timeout) @@ -320,13 +353,14 @@ def _crawl(self, "chunks_per_source": chunks_per_source, } + override_headers = self._pop_request_headers(kwargs) if kwargs: data.update(kwargs) - + data = {k: v for k, v in data.items() if v is not None} try: - response = self.session.post(self.base_url + "/crawl", data=json.dumps(data), timeout=timeout) + response = self.session.post(self.base_url + "/crawl", data=json.dumps(data), timeout=timeout, **({"headers": override_headers} if override_headers else {})) except requests.exceptions.Timeout: raise TimeoutError(timeout) @@ -428,13 +462,14 @@ def _map(self, "include_usage": include_usage, } + override_headers = self._pop_request_headers(kwargs) if kwargs: data.update(kwargs) - + data = {k: v for k, v in data.items() if v is not None} try: - response = self.session.post(self.base_url + "/map", data=json.dumps(data), timeout=timeout) + response = self.session.post(self.base_url + "/map", data=json.dumps(data), timeout=timeout, **({"headers": override_headers} if override_headers else {})) except requests.exceptions.Timeout: raise TimeoutError(timeout) @@ -595,6 +630,7 @@ def _research(self, data = {k: v for k, v in data.items() if v is not None} + override_headers = self._pop_request_headers(kwargs) if kwargs: data.update(kwargs) @@ -604,7 +640,8 @@ def _research(self, self.base_url + "/research", data=json.dumps(data), timeout=timeout, - stream=True + stream=True, + **({"headers": override_headers} if override_headers else {}) ) except requests.exceptions.Timeout: raise TimeoutError(timeout) @@ -641,7 +678,8 @@ def stream_generator() -> Generator[bytes, None, None]: response = self.session.post( self.base_url + "/research", data=json.dumps(data), - timeout=timeout + timeout=timeout, + **({"headers": override_headers} if override_headers else {}) ) except requests.exceptions.Timeout: raise TimeoutError(timeout)