From 31d0694c25470b36814f0e767494b95b5c3cc47f Mon Sep 17 00:00:00 2001
From: tinosattavily <konstantinos.a@tavily.com>
Date: Mon, 26 Jan 2026 13:27:59 -0500
Subject: [PATCH] perf(client): add session pooling for connection reuse

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
---
 tavily/tavily.py | 31 +++++++++++++------------------
 1 file changed, 13 insertions(+), 18 deletions(-)

diff --git a/tavily/tavily.py b/tavily/tavily.py
index abd22ba..ad96dc5 100644
--- a/tavily/tavily.py
+++ b/tavily/tavily.py
@@ -38,6 +38,11 @@ def __init__(self, api_key: Optional[str] = None, proxies: Optional[dict[str, st
             **({"X-Project-ID": tavily_project} if tavily_project else {})
         }
 
+        self.session = requests.Session()
+        self.session.headers.update(self.headers)
+        if self.proxies:
+            self.session.proxies.update(self.proxies)
+
     def _search(self,
                 query: str,
                 search_depth: Literal["basic", "advanced", "fast", "ultra-fast"] = None,
@@ -91,7 +96,7 @@ def _search(self,
         timeout = min(timeout, 120)
 
         try:
-            response = requests.post(self.base_url + "/search", data=json.dumps(data), headers=self.headers, timeout=timeout, proxies=self.proxies)
+            response = self.session.post(self.base_url + "/search", data=json.dumps(data), timeout=timeout)
         except requests.exceptions.Timeout:
             raise TimeoutError(timeout)
 
@@ -202,7 +207,7 @@ def _extract(self,
             data.update(kwargs)
 
         try:
-            response = requests.post(self.base_url + "/extract", data=json.dumps(data), headers=self.headers, timeout=timeout, proxies=self.proxies)
+            response = self.session.post(self.base_url + "/extract", data=json.dumps(data), timeout=timeout)
         except requests.exceptions.Timeout:
             raise TimeoutError(timeout)
 
@@ -310,8 +315,7 @@ def _crawl(self,
         data = {k: v for k, v in data.items() if v is not None}
 
         try:
-            response = requests.post(
-                self.base_url + "/crawl", data=json.dumps(data), headers=self.headers, timeout=timeout, proxies=self.proxies)
+            response = self.session.post(self.base_url + "/crawl", data=json.dumps(data), timeout=timeout)
         except requests.exceptions.Timeout:
             raise TimeoutError(timeout)
 
@@ -421,8 +425,7 @@ def _map(self,
         data = {k: v for k, v in data.items() if v is not None}
 
         try:
-            response = requests.post(
-                self.base_url + "/map", data=json.dumps(data), headers=self.headers, timeout=timeout, proxies=self.proxies)
+            response = self.session.post(self.base_url + "/map", data=json.dumps(data), timeout=timeout)
         except requests.exceptions.Timeout:
             raise TimeoutError(timeout)
 
@@ -631,12 +634,10 @@ def _research(self,
 
         if stream:
             try:
-                response = requests.post(
+                response = self.session.post(
                     self.base_url + "/research",
                     data=json.dumps(data),
-                    headers=self.headers,
                     timeout=timeout,
-                    proxies=self.proxies,
                     stream=True
                 )
             except requests.exceptions.Timeout:
@@ -671,12 +672,10 @@ def stream_generator() -> Generator[bytes, None, None]:
             return stream_generator()
         else:
             try:
-                response = requests.post(
+                response = self.session.post(
                     self.base_url + "/research",
                     data=json.dumps(data),
-                    headers=self.headers,
-                    timeout=timeout,
-                    proxies=self.proxies
+                    timeout=timeout
                 )
             except requests.exceptions.Timeout:
                 raise TimeoutError(timeout)
@@ -752,11 +751,7 @@ def get_research(self,
             dict: Research response containing request_id, created_at, completed_at, status, content, and sources.
         """
         try:
-            response = requests.get(
-                self.base_url + f"/research/{request_id}",
-                headers=self.headers,
-                proxies=self.proxies,
-            )
+            response = self.session.get(self.base_url + f"/research/{request_id}")
         except Exception as e:
             raise Exception(f"Error getting research: {e}")