Merge pull request #791 from parea-ai/PAI-1069-add-sdk-header

feat: integrate with sglang

Author: joschkabraun

parea/__init__.py

@@ -9,8 +9,8 @@
     run the following command:  ```bash pip install parea ```.
 """
 import sys
-from importlib import metadata as importlib_metadata
 
+from parea.api_client import get_version
 from parea.cache import InMemoryCache
 from parea.client import Parea
 from parea.experiment.cli import experiment as _experiment_cli
@@ -22,14 +22,6 @@
 from parea.wrapper.openai_raw_api_tracer import aprocess_stream_and_yield, process_stream_and_yield
 from parea.wrapper.utils import convert_openai_raw_to_log
 
-
-def get_version() -> str:
-    try:
-        return importlib_metadata.version(__name__)
-    except importlib_metadata.PackageNotFoundError:  # pragma: no cover
-        return "unknown"
-
-
 version: str = get_version()
 
 

parea/api_client.py

@@ -1,10 +1,11 @@
-from typing import Any, AsyncIterable, Callable, Dict, Optional
+from typing import Any, AsyncIterable, Callable, Dict, List, Optional
 
 import asyncio
 import json
 import os
 import time
 from functools import wraps
+from importlib import metadata as importlib_metadata
 
 import httpx
 from dotenv import load_dotenv
@@ -60,6 +61,7 @@ class HTTPClient:
     _instance = None
     base_url = os.getenv("PAREA_BASE_URL", "https://parea-ai-backend-us-9ac16cdbc7a7b006.onporter.run/api/parea/v1")
     api_key = None
+    integrations: List[str] = []
 
     def __new__(cls, *args, **kwargs):
         if cls._instance is None:
@@ -71,6 +73,20 @@ def __new__(cls, *args, **kwargs):
     def set_api_key(self, api_key: str):
         self.api_key = api_key
 
+    def add_integration(self, integration: str):
+        if integration not in self.integrations:
+            self.integrations.append(integration)
+
+    def _get_headers(self, api_key: Optional[str] = None) -> Dict[str, str]:
+        headers = {
+            "x-api-key": self.api_key or api_key,
+            "x-sdk-version": get_version(),
+            "x-sdk-language": "python",
+        }
+        if self.integrations:
+            headers["x-sdk-integrations"] = ",".join(self.integrations)
+        return headers
+
     @retry_on_502
     def request(
         self,
@@ -83,7 +99,7 @@ def request(
         """
         Makes an HTTP request to the specified endpoint.
         """
-        headers = {"x-api-key": self.api_key} if self.api_key else api_key
+        headers = self._get_headers(api_key=api_key)
         try:
             response = self.sync_client.request(method, endpoint, json=data, headers=headers, params=params)
             response.raise_for_status()
@@ -106,7 +122,7 @@ async def request_async(
         """
         Makes an asynchronous HTTP request to the specified endpoint.
         """
-        headers = {"x-api-key": self.api_key} if self.api_key else api_key
+        headers = self._get_headers(api_key=api_key)
         try:
             response = await self.async_client.request(method, endpoint, json=data, headers=headers, params=params)
             response.raise_for_status()
@@ -128,7 +144,7 @@ def stream_request(
         """
         Makes a streaming HTTP request to the specified endpoint, yielding chunks of data.
         """
-        headers = {"x-api-key": self.api_key} if self.api_key else api_key
+        headers = self._get_headers(api_key=api_key)
         try:
             with self.sync_client.stream(method, endpoint, json=data, headers=headers, params=params, timeout=None) as response:
                 response.raise_for_status()
@@ -151,7 +167,7 @@ async def stream_request_async(
         """
         Makes an asynchronous streaming HTTP request to the specified endpoint, yielding chunks of data.
         """
-        headers = {"x-api-key": self.api_key} if self.api_key else api_key
+        headers = self._get_headers(api_key=api_key)
         try:
             async with self.async_client.stream(method, endpoint, json=data, headers=headers, params=params, timeout=None) as response:
                 response.raise_for_status()
@@ -200,3 +216,10 @@ def parse_event_data(byte_data):
     except Exception as e:
         print(f"Error parsing event data: {e}")
         return None
+
+
+def get_version() -> str:
+    try:
+        return importlib_metadata.version("parea-ai")
+    except importlib_metadata.PackageNotFoundError:  # pragma: no cover
+        return "unknown"

parea/client.py

@@ -76,19 +76,25 @@ def __attrs_post_init__(self):
             parea_logger.set_client(self._client)
             parea_logger.set_project_uuid(self.project_uuid)
 
-    def wrap_openai_client(self, client: "OpenAI") -> None:
+    def wrap_openai_client(self, client: "OpenAI", integration: Optional[str] = None) -> None:
         """Only necessary for instance client with OpenAI version >= 1.0.0"""
         from parea.wrapper import OpenAIWrapper
         from parea.wrapper.openai_beta_wrapper import BetaWrappers
 
         OpenAIWrapper().init(log=logger_all_possible, cache=self.cache, module_client=client)
         BetaWrappers(client).init()
 
-    def wrap_anthropic_client(self, client: "Anthropic") -> None:
+        if integration:
+            self._client.add_integration(integration)
+
+    def wrap_anthropic_client(self, client: "Anthropic", integration: Optional[str] = None) -> None:
         from parea.wrapper.anthropic.anthropic import AnthropicWrapper
 
         AnthropicWrapper().init(log=logger_all_possible, cache=self.cache, client=client)
 
+        if integration:
+            self._client.add_integration(integration)
+
     def auto_trace_openai_clients(self) -> None:
         import openai
 
@@ -97,6 +103,10 @@ def auto_trace_openai_clients(self) -> None:
         openai.AzureOpenAI = patch_openai_client_classes(openai.AzureOpenAI, self)
         openai.AsyncAzureOpenAI = patch_openai_client_classes(openai.AsyncAzureOpenAI, self)
 
+    def integrate_with_sglang(self):
+        self.auto_trace_openai_clients()
+        self._client.add_integration("sglang")
+
     def _add_project_uuid_to_data(self, data) -> dict:
         data_dict = asdict(data)
         data_dict["project_uuid"] = self._project.uuid

parea/parea_logger.py

@@ -66,6 +66,7 @@ def record_vendor_log(self, data: Dict[str, Any], vendor: TraceIntegrations) ->
         data["project_uuid"] = self._project_uuid
         if experiment_uuid := os.getenv(PAREA_OS_ENV_EXPERIMENT_UUID, None):
             data["experiment_uuid"] = experiment_uuid
+        self._client.add_integration("langchain")
         self._client.request(
             "POST",
             VENDOR_LOG_ENDPOINT.format(vendor=vendor.value),
@@ -76,6 +77,7 @@ async def arecord_vendor_log(self, data: Dict[str, Any], vendor: TraceIntegratio
         data["project_uuid"] = self._project_uuid
         if experiment_uuid := os.getenv(PAREA_OS_ENV_EXPERIMENT_UUID, None):
             data["experiment_uuid"] = experiment_uuid
+        self._client.add_integration("langchain")
         await self._client.request_async(
             "POST",
             VENDOR_LOG_ENDPOINT.format(vendor=vendor.value),

pyproject.toml

@@ -6,7 +6,7 @@ build-backend = "poetry.core.masonry.api"
 [tool.poetry]
 name = "parea-ai"
 packages = [{ include = "parea" }]
-version = "0.2.138"
+version = "0.2.139"
 description = "Parea python sdk"
 readme = "README.md"
 authors = ["joel-parea-ai <[email protected]>"]