Merge branch 'main' into PAI-449-improve-sdk-trace-log-system-architecture-pys

Author: jalexanderII

.github/workflows/build.yml

@@ -7,7 +7,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: [ "3.9" ]
+        python-version: [ "3.8" ]
 
     steps:
       - uses: actions/[email protected]
@@ -33,9 +33,9 @@ jobs:
         run: |
           make check-codestyle
 
-#      - name: Run tests
-#        run: |
-#          make test
+      - name: Run tests
+        run: |
+          make test
 
 #    - name: Run safety checks
 #      run: |

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,17 +1,18 @@
-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
 
 load_dotenv()
 
-MAX_RETRIES = 7
+MAX_RETRIES = 8
 BACKOFF_FACTOR = 0.5
 
 
@@ -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

@@ -14,7 +14,7 @@
 from parea.cache.cache import Cache
 from parea.constants import PAREA_OS_ENV_EXPERIMENT_UUID
 from parea.experiment.datasets import create_test_cases, create_test_collection
-from parea.helpers import gen_trace_id, serialize_metadata_values, structure_trace_log_from_api
+from parea.helpers import gen_trace_id, serialize_metadata_values, structure_trace_log_from_api, structure_trace_logs_from_api
 from parea.parea_logger import parea_logger
 from parea.schemas.models import (
     Completion,
@@ -25,11 +25,14 @@
     CreateTestCases,
     ExperimentSchema,
     ExperimentStatsSchema,
+    ExperimentWithPinnedStatsSchema,
     FeedbackRequest,
     FinishExperimentRequestSchema,
+    ListExperimentUUIDsFilters,
     ProjectSchema,
     TestCaseCollection,
     TraceLog,
+    TraceLogFilters,
     UseDeployedPrompt,
     UseDeployedPromptResponse,
 )
@@ -50,6 +53,8 @@
 CREATE_COLLECTION_ENDPOINT = "/collection"
 ADD_TEST_CASES_ENDPOINT = "/testcases"
 GET_TRACE_LOG_ENDPOINT = "/trace_log/{trace_id}"
+LIST_EXPERIMENTS_ENDPOINT = "/experiments"
+GET_EXPERIMENT_LOGS_ENDPOINT = "/experiment/{experiment_uuid}/trace_logs"
 
 
 @define
@@ -71,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
 
@@ -92,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
@@ -346,6 +361,22 @@ async def aget_trace_log(self, trace_id: str) -> TraceLog:
         response = await self._client.request_async("GET", GET_TRACE_LOG_ENDPOINT.format(trace_id=trace_id))
         return structure_trace_log_from_api(response.json())
 
+    def list_experiments(self, filter_conditions: Optional[ListExperimentUUIDsFilters] = ListExperimentUUIDsFilters()) -> List[ExperimentWithPinnedStatsSchema]:
+        response = self._client.request("POST", LIST_EXPERIMENTS_ENDPOINT, data=asdict(filter_conditions))
+        return structure(response.json(), List[ExperimentWithPinnedStatsSchema])
+
+    async def alist_experiments(self, filter_conditions: Optional[ListExperimentUUIDsFilters] = ListExperimentUUIDsFilters()) -> List[ExperimentWithPinnedStatsSchema]:
+        response = await self._client.request_async("POST", LIST_EXPERIMENTS_ENDPOINT, data=asdict(filter_conditions))
+        return structure(response.json(), List[ExperimentWithPinnedStatsSchema])
+
+    def get_experiment_trace_logs(self, experiment_uuid: str, filters: TraceLogFilters = TraceLogFilters()) -> List[TraceLog]:
+        response = self._client.request("POST", GET_EXPERIMENT_LOGS_ENDPOINT.format(experiment_uuid=experiment_uuid), data=asdict(filters))
+        return structure_trace_logs_from_api(response.json())
+
+    async def aget_experiment_trace_logs(self, experiment_uuid: str, filters: TraceLogFilters = TraceLogFilters()) -> List[TraceLog]:
+        response = await self._client.request_async("POST", GET_EXPERIMENT_LOGS_ENDPOINT.format(experiment_uuid=experiment_uuid), data=asdict(filters))
+        return structure_trace_logs_from_api(response.json())
+
 
 _initialized_parea_wrapper = False
 

parea/cookbook/list_experiments.py

@@ -0,0 +1,16 @@
+import os
+
+from dotenv import load_dotenv
+
+from parea import Parea
+from parea.schemas import ListExperimentUUIDsFilters
+
+load_dotenv()
+
+p = Parea(api_key=os.getenv("PAREA_API_KEY"))
+
+experiments = p.list_experiments(ListExperimentUUIDsFilters(experiment_name_filter="Greeting"))
+print(f"Num. experiments: {len(experiments)}")
+trace_logs = p.get_experiment_trace_logs(experiments[0].uuid)
+print(f"Num. trace logs: {len(trace_logs)}")
+print(f"Trace log: {trace_logs[0]}")

parea/cookbook/tracing_with_images_anthropic.py

@@ -0,0 +1,74 @@
+import json
+import os
+
+from anthropic import Anthropic
+from dotenv import load_dotenv
+from openai import OpenAI
+
+from parea import Parea, trace, trace_insert
+from parea.schemas import TraceLogImage
+
+load_dotenv()
+
+
+oai_client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))
+a_client = Anthropic(api_key=os.getenv("ANTHROPIC_API_KEY"))
+
+p = Parea(api_key=os.getenv("PAREA_API_KEY"))
+p.wrap_openai_client(oai_client)
+p.wrap_anthropic_client(a_client)
+
+
+@trace
+def image_maker(query: str) -> str:
+    response = oai_client.images.generate(prompt=query, model="dall-e-3")
+    image_url = response.data[0].url
+    caption = {"original_prompt": query, "revised_prompt": response.data[0].revised_prompt}
+    trace_insert({"images": [TraceLogImage(url=image_url, caption=json.dumps(caption))]})
+    return image_url
+
+
+from typing import Optional
+
+import base64
+
+import requests
+
+
+@trace
+def ask_vision(image_url: str) -> Optional[str]:
+    image_data = requests.get(image_url).content
+    base64_image = base64.b64encode(image_data).decode("utf-8")
+
+    response = a_client.messages.create(
+        model="claude-3-haiku-20240307",
+        messages=[
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "image",
+                        "source": {
+                            "type": "base64",
+                            "media_type": "image/png",
+                            "data": base64_image,
+                        },
+                    },
+                    {"type": "text", "text": "What’s in this image?"},
+                ],
+            }
+        ],
+        max_tokens=300,
+    )
+    return response.content[0].text
+
+
+@trace
+def main(query: str) -> str:
+    image_url = image_maker(query)
+    return ask_vision(image_url)
+
+
+if __name__ == "__main__":
+    result = main("A dog sitting comfortably on a chair")
+    print(result)

parea/cookbook/tracing_with_images_open_ai.py

@@ -30,7 +30,7 @@ def image_maker(query: str) -> str:
 @trace
 def ask_vision(image_url: str) -> Optional[str]:
     response = client.chat.completions.create(
-        model="gpt-4-vision-preview",
+        model="gpt-4-turbo",
         messages=[
             {
                 "role": "user",

parea/evals/rag/context_query_relevancy.py

@@ -1,6 +1,6 @@
 from typing import Callable, List, Optional
 
-from parea.evals.utils import call_openai, sent_tokenize
+from parea.evals.utils import call_openai, get_context, sent_tokenize
 from parea.schemas.log import Log
 
 
@@ -27,13 +27,7 @@ def context_query_relevancy_factory(
     def context_query_relevancy(log: Log) -> float:
         """Quantifies how much the retrieved context relates to the query."""
         question = log.inputs[question_field]
-        if context_fields:
-            context = "\n".join(log.inputs[context_field] for context_field in context_fields)
-        else:
-            if isinstance(log.output, list):
-                context = "\n".join(log.output)
-            else:
-                context = str(log.output)
+        context = get_context(log, context_fields)
 
         extracted_sentences = call_openai(
             model=model,

parea/evals/rag/context_ranking_listwise.py

@@ -1,6 +1,6 @@
 from typing import Callable, List, Optional
 
-from parea.evals.utils import call_openai, ndcg
+from parea.evals.utils import call_openai, get_context, ndcg
 from parea.schemas.log import Log
 
 
@@ -99,13 +99,7 @@ def progressive_reranking(query: str, contexts: List[str]) -> List[int]:
     def context_ranking(log: Log) -> float:
         """Quantifies if the retrieved context is ranked by their relevancy by re-ranking the contexts."""
         question = log.inputs[question_field]
-        if context_fields:
-            contexts = [log.inputs[context_field] for context_field in context_fields]
-        else:
-            if isinstance(log.output, list):
-                contexts = log.output
-            else:
-                contexts = [str(log.output)]
+        contexts = get_context(log, context_fields, True)
 
         reranked_indices = progressive_reranking(question, contexts)