Merge pull request #743 from parea-ai/PAI-1006-add-session-id-and-filter-by-session-id-p

feat(log): add session_id

Author: jalexanderII

parea/cookbook/tracing_anthropic.py

@@ -13,7 +13,7 @@
 client = anthropic.Anthropic()
 aclient = anthropic.AsyncAnthropic()
 
-p = Parea(api_key=os.getenv("PAREA_API_KEY"), project_name="testing")
+p = Parea(api_key=os.getenv("PAREA_API_KEY"))
 p.wrap_anthropic_client(client)
 p.wrap_anthropic_client(aclient)
 

parea/cookbook/tracing_with_open_ai_endpoint_directly.py

@@ -6,7 +6,7 @@
 from dotenv import load_dotenv
 from openai import OpenAI
 
-from parea import Parea, get_current_trace_id, trace
+from parea import Parea, get_current_trace_id, trace, trace_insert
 from parea.schemas import FeedbackRequest
 
 load_dotenv()
@@ -73,13 +73,14 @@ def refiner(query: str, additional_description: str, argument: str, criticism: s
 @trace
 def argument_chain(query: str, additional_description: str = "") -> Tuple[str, str]:
     trace_id = get_current_trace_id()
+    trace_insert({"session_id": "cus_1234", "end_user_identifier": "user_1234"}, trace_id)
     argument = argumentor(query, additional_description)
     criticism = critic(argument)
     refined_argument = refiner(query, additional_description, argument, criticism)
     return refined_argument, trace_id
 
 
-@trace
+@trace(session_id="cus_1234", end_user_identifier="user_1234")
 def json_call() -> str:
     completion = client.chat.completions.create(
         model="gpt-4-turbo-2024-04-09",
@@ -95,11 +96,11 @@ def json_call() -> str:
         additional_description="Provide a concise, few sentence argument on why sparkling wine is good for you.",
     )
     print(result)
-    # p.record_feedback(
-    #     FeedbackRequest(
-    #         trace_id=trace_id,
-    #         score=0.7,  # 0.0 (bad) to 1.0 (good)
-    #     )
-    # )
+    p.record_feedback(
+        FeedbackRequest(
+            trace_id=trace_id,
+            score=0.7,  # 0.0 (bad) to 1.0 (good)
+        )
+    )
 
     print(json_call())

parea/schemas/log.py

@@ -43,67 +43,7 @@ class ModelParams:
 
 @define
 class LLMInputs:
-    """
-    model choice should match an enabled model on the Parea platform.
-    model options are:
-    OPENAI_MODELS
-        "gpt-3.5-turbo"
-        "gpt-3.5-turbo-0301"
-        "gpt-3.5-turbo-0613"
-        "gpt-3.5-turbo-16k"
-        "gpt-3.5-turbo-16k-0301"
-        "gpt-3.5-turbo-16k-0613"
-        "gpt-3.5-turbo-1106"
-        "gpt-3.5-turbo-0125"
-        "gpt-3.5-turbo-instruct"
-        "gpt-4"
-        "gpt-4-0314"
-        "gpt-4-0613"
-        "gpt-4-32k"
-        "gpt-4-32k-0314"
-        "gpt-4-32k-0613"
-        "gpt-4-turbo-preview"
-        "gpt-4-1106-preview"
-        "gpt-4-0125-preview"
-    You can use Azure models by providing your model name prefixed with azure_, e.g. azure_gpt-3.5-turbo
-    ANTHROPIC_MODELS
-        "claude-instant-1.1"
-        "claude-instant-1"
-        "claude-instant-1.2"
-        "claude-instant-1-100k"
-        "claude-instant-1.1-100k"
-        "claude-1"
-        "claude-2"
-        "claude-1-100k"
-        "claude-1.2"
-        "claude-1.3"
-        "claude-1.3-100k"
-        "claude-2.1"
-    AWS_ANTHROPIC_MODELS
-        "anthropic.claude-instant-v1"
-        "anthropic.claude-v1"
-        "anthropic.claude-v2"
-        "anthropic.claude-v2:1"
-    ANYSCALE_MODELS
-        "meta-llama/Llama-2-7b-chat-hf"
-        "meta-llama/Llama-2-13b-chat-hf"
-        "meta-llama/Llama-2-70b-chat-hf"
-        "codellama/CodeLlama-34b-Instruct-hf"
-        "mistralai/Mistral-7B-Instruct-v0.1"
-        "mistralai/Mixtral-8x7B-Instruct-v0.1"
-        "HuggingFaceH4/zephyr-7b-beta"
-        "Open-Orca/Mistral-7B-OpenOrca"
-    VERTEX_MODELS
-        "gemini-pro"
-        "text-bison@001"
-        "text-bison@002"
-        "text-bison"
-        "text-bison-32k"
-        "chat-bison@002"
-        "chat-bison@001"
-        "chat-bison"
-        "chat-bison-32k
-    """
+    """model choice should match an enabled model on the Parea platform."""
 
     model: Optional[str] = None
     provider: Optional[str] = None

parea/schemas/models.py

@@ -114,6 +114,7 @@ class TraceLog(EvaluatedLog):
     # metrics filled from either decorator or completion
     end_timestamp: Optional[str] = None
     end_user_identifier: Optional[str] = None
+    session_id: Optional[str] = None
     metadata: Optional[Dict[str, Any]] = None
     tags: Optional[List[str]] = field(factory=list)
     experiment_uuid: Optional[str] = None

parea/utils/trace_utils.py

@@ -89,7 +89,7 @@ def trace_insert(data: Dict[str, Any], trace_id: Optional[str] = None):
     Insert data into the trace log for the current or specified trace id. Data should be a dictionary with keys that correspond to the fields of the TraceLog model.
     If the field already has an existing value that is extensible (dict, set, list, etc.), the new value will be merged with the existing value.
     Args:
-        data: Keys can be one of: trace_name, end_user_identifier, metadata, tags, deployment_id, images
+        data: Keys can be one of: trace_name, end_user_identifier, metadata, tags, deployment_id, images, session_id
         trace_id: The trace id to insert the data into. If not provided, the current trace id will be used.
     """
     try:
@@ -128,6 +128,7 @@ def trace(
     tags: Optional[List[str]] = None,
     metadata: Optional[Dict[str, Any]] = None,
     end_user_identifier: Optional[str] = None,
+    session_id: Optional[str] = None,
     eval_funcs_names: Optional[List[str]] = None,
     eval_funcs: Optional[List[Callable]] = None,
     access_output_of_func: Optional[Callable] = None,
@@ -171,6 +172,7 @@ def init_trace(func_name, _parea_target_field, args, kwargs, func) -> Tuple[str,
                 start_timestamp=start_time.isoformat(),
                 trace_name=name or func_name,
                 end_user_identifier=end_user_identifier,
+                session_id=session_id,
                 metadata=metadata,
                 target=_parea_target_field,
                 tags=tags,

parea/wrapper/wrapper.py

@@ -89,6 +89,7 @@ def _init_trace(self) -> Tuple[str, datetime, contextvars.Token]:
                 start_timestamp=start_time.isoformat(),
                 trace_name="LLM",
                 end_user_identifier=None,
+                session_id=None,
                 metadata=None,
                 target=None,
                 tags=None,