Merge pull request #17 from weavel-ai/aschung01/wea-509-fix-set-generation-name-as-optional-param

fix: set generation name as optional param

Author: engineerA314

README.md

@@ -23,7 +23,7 @@ You can find our full documentation [here](https://weavel.ai/docs/python-sdk).
 
 ## How to use
 
-### Basic Usage
+### Option 1: Using OpenAI wrapper
 
 ```python
 from weavel import WeavelOpenAI as OpenAI
@@ -42,7 +42,40 @@ response = openai.chat.completions.create(
 
 ```
 
-### Advanced Usage
+### Option 2: Logging inputs/outputs of LLM calls
+
+```python
+from weavel import Weavel
+from openai import OpenAI
+from pydantic import BaseModel
+
+openai = OpenAI()
+# initialize Weavel
+weavel = Weavel()
+
+class Answer(BaseModel):
+    reasoning: str
+    answer: str
+
+question = "What is x if x + 2 = 4?"
+response = openai.beta.chat.completions.parse(
+    model="gpt-4o-2024-08-06",
+    messages=[
+        {"role": "system", "content": "You are a math teacher."},
+        {"role": "user", "content": question}
+    ],
+    response_format=Answer
+).choices[0].message.parsed
+
+# log the generation
+weavel.generation(
+    name="solve-math", # optional
+    inputs={"question": question},
+    outputs=response.model_dump()
+)
+```
+
+### Option 3 (Advanced Usage): OTEL-compatible trace logging
 
 ```python
 from weavel import Weavel

setup.py

@@ -10,7 +10,7 @@
 
 setup(
     name="weavel",
-    version="1.7.1",
+    version="1.8.0",
     packages=find_namespace_packages(),
     entry_points={},
     description="Weavel, Prompt Optimization and Evaluation for LLM Applications",
@@ -32,8 +32,6 @@
         "termcolor",
         "watchdog",
         "readerwriterlock",
-        "pendulum",
-        "httpx[http2]",
         "nest_asyncio",
     ],
     python_requires=">=3.8.10",
@@ -47,6 +45,6 @@
         "dataset curation",
         "prompt engineering",
         "prompt optimization",
-        "AI Prompt Engineer"
+        "AI Prompt Engineer",
     ],
 )

weavel/_body.py

@@ -132,7 +132,10 @@ class CaptureRecordBody(BaseRecordBody):
 
 
 class CaptureObservationBody(BaseObservationBody):
-    name: str
+    name: Optional[str] = Field(
+        default=None,
+        description="The name of the observation. Optional.",
+    )
 
 
 class CaptureMessageBody(CaptureRecordBody):
@@ -262,6 +265,22 @@ class CaptureGenerationBody(CaptureObservationBody):
         default=None,
         description="The outputs of the generation. Optional.",
     )
+    messages: Optional[List[Dict[str, str]]] = Field(
+        default=None,
+        description="The messages of the generation. Optional.",
+    )
+    model: Optional[str] = Field(
+        default=None,
+        description="The model of the generation. Optional.",
+    )
+    latency: Optional[float] = Field(
+        default=None,
+        description="The latency of the generation. Optional.",
+    )
+    cost: Optional[float] = Field(
+        default=None,
+        description="The cost of the generation. Optional.",
+    )
     prompt_name: Optional[str] = Field(
         default=None,
         description="The name of the prompt. Optional.",

weavel/_worker.py

@@ -341,11 +341,15 @@ def capture_generation(
         self,
         observation_id: str,
         created_at: datetime,
-        name: str,
+        name: Optional[str] = None,
         record_id: Optional[str] = None,
         inputs: Optional[Union[Dict[str, Any], List[Any], str]] = None,
         outputs: Optional[Union[Dict[str, Any], List[Any], str]] = None,
+        messages: Optional[List[Dict[str, str]]] = None,
         metadata: Optional[Dict[str, str]] = None,
+        model: Optional[str] = None,
+        latency: Optional[float] = None,
+        cost: Optional[float] = None,
         parent_observation_id: Optional[str] = None,
         prompt_name: Optional[str] = None,
     ):
@@ -358,6 +362,10 @@ def capture_generation(
                     name=name,
                     inputs=inputs,
                     outputs=outputs,
+                    messages=messages,
+                    model=model,
+                    latency=latency,
+                    cost=cost,
                     metadata=metadata,
                     parent_observation_id=parent_observation_id,
                     prompt_name=prompt_name,
@@ -374,6 +382,7 @@ async def acapture_generation(
         record_id: Optional[str] = None,
         inputs: Optional[Union[Dict[str, Any], List[Any], str]] = None,
         outputs: Optional[Union[Dict[str, Any], List[Any], str]] = None,
+        messages: Optional[List[Dict[str, str]]] = None,
         metadata: Optional[Dict[str, str]] = None,
         parent_observation_id: Optional[str] = None,
         prompt_name: Optional[str] = None,
@@ -387,6 +396,7 @@ async def acapture_generation(
                     name=name,
                     inputs=inputs,
                     outputs=outputs,
+                    messages=messages,
                     metadata=metadata,
                     parent_observation_id=parent_observation_id,
                     prompt_name=prompt_name,

weavel/client.py

@@ -190,6 +190,7 @@ def generation(
         name: Optional[str] = None,
         inputs: Optional[Union[Dict[str, Any], List[Any], str]] = None,
         outputs: Optional[Union[Dict[str, Any], List[Any], str]] = None,
+        messages: Optional[List[Dict[str, str]]] = None,
         metadata: Optional[Dict[str, Any]] = None,
         parent_observation_id: Optional[str] = None,
         prompt_name: Optional[str] = None,
@@ -212,16 +213,11 @@ def generation(
                 name=name,
                 inputs=inputs,
                 outputs=outputs,
+                messages=messages,
                 metadata=metadata,
                 parent_observation_id=parent_observation_id,
                 prompt_name=prompt_name,
             )
-        if (
-            record_id or (record_id is None and parent_observation_id)
-        ) and name is None:
-            raise ValueError(
-                "If you want to create a new generation, you must provide a name."
-            )
 
         if record_id is not None or (
             record_id is None and parent_observation_id is not None
@@ -240,6 +236,7 @@ def generation(
                     name=name,
                     inputs=inputs,
                     outputs=outputs,
+                    messages=messages,
                     metadata=metadata,
                     parent_observation_id=parent_observation_id,
                     prompt_name=prompt_name,
@@ -251,6 +248,7 @@ def generation(
                 name=name,
                 inputs=inputs,
                 outputs=outputs,
+                messages=messages,
                 metadata=metadata,
                 parent_observation_id=parent_observation_id,
                 prompt_name=prompt_name,
@@ -442,8 +440,8 @@ def create_dataset_items(
             return
 
         for item in items:
-            if not isinstance(item, DatasetItem):
-                item = DatasetItem(**item)
+            if isinstance(item, DatasetItem):
+                item = item.model_dump()
 
         self._worker.create_dataset_items(dataset_name, items)
 
@@ -462,8 +460,8 @@ async def acreate_dataset_items(
             return
 
         for item in items:
-            if not isinstance(item, DatasetItem):
-                item = DatasetItem(**item)
+            if isinstance(item, DatasetItem):
+                item = item.model_dump()
 
         await self._worker.acreate_dataset_items(dataset_name, items)
 

weavel/wrapper.py

@@ -41,6 +41,7 @@
 
 load_dotenv()
 
+
 def calculate_cost(model: str, usage: Dict[str, int]) -> float:
     if model not in pricing:
         # raise ValueError(f"Unknown model: {model}")
@@ -240,6 +241,7 @@ def _handle_non_streaming(self, *args, **kwargs):
                 cost = calculate_cost(response.model, formatted_response["usage"])
                 formatted_response["cost"] = cost
 
+                # TODO: models, latency, cost shouldn't be logged as inputs. Should have separate fields for each.
                 self._capture_generation(inputs=kwargs, outputs=formatted_response)
                 return response
 
@@ -275,11 +277,11 @@ def parse(self, *args, **kwargs):
                 # context manager to suppress duplicate logging
                 with LoggingSuppressor(self._worker):
                     response = self.original.parse(*args, **kwargs)
-                    
+
                 end_time = time.time()
                 latency = end_time - start_time
                 latency = round(latency, 6)
-                
+
                 formatted_response = {
                     "id": response.id,
                     "choices": [choice.model_dump() for choice in response.choices],
@@ -291,11 +293,12 @@ def parse(self, *args, **kwargs):
                     "system_fingerprint": response.system_fingerprint,
                     "latency": latency,
                 }
-                
+
                 # Calculate and add cost
                 cost = calculate_cost(response.model, formatted_response["usage"])
                 formatted_response["cost"] = cost
-                
+
+                # TODO: models, latency, cost shouldn't be logged as inputs. Should have separate fields for each.
                 self._worker.capture_generation(
                     observation_id=str(uuid4()),
                     created_at=datetime.now(timezone.utc),
@@ -304,7 +307,7 @@ def parse(self, *args, **kwargs):
                     inputs=kwargs,
                     outputs=formatted_response,
                 )
-                
+
                 return response
 
         self.chat.completions = CustomChatCompletions(
@@ -442,7 +445,7 @@ async def _handle_streaming(self, *args, **kwargs):
                 usage = accumulated_response["usage"]
                 cost = calculate_cost(model, usage)
                 accumulated_response["cost"] = cost
-                
+
                 if self._worker.capture_generation is not None:
                     await self._capture_generation(
                         inputs=kwargs, outputs=accumulated_response
@@ -478,7 +481,7 @@ async def _handle_non_streaming(self, *args, **kwargs):
                     inputs=kwargs, outputs=formatted_response
                 )
                 return response
-            
+
             async def _capture_generation(
                 self,
                 inputs,
@@ -507,14 +510,14 @@ def __init__(self, original, weavel_api_key: str, base_url: Optional[str]):
             async def parse(self, *args, **kwargs):
                 header = kwargs.pop("headers", {})
                 start_time = time.time()
-                
+
                 async with AsyncLoggingSuppressor(self._worker):
                     response = await self.original.parse(*args, **kwargs)
-                    
+
                 end_time = time.time()
                 latency = end_time - start_time
                 latency = round(latency, 6)
-                    
+
                 formatted_response = {
                     "id": response.id,
                     "choices": [choice.model_dump() for choice in response.choices],
@@ -526,7 +529,7 @@ async def parse(self, *args, **kwargs):
                     "system_fingerprint": response.system_fingerprint,
                     "latency": latency,
                 }
-                
+
                 # Calculate and add cost
                 cost = calculate_cost(response.model, formatted_response["usage"])
                 formatted_response["cost"] = cost