airtrain 0.1.27__tar.gz → 0.1.30__tar.gz

{airtrain-0.1.27/airtrain.egg-info → airtrain-0.1.30}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: airtrain
-Version: 0.1.27
+Version: 0.1.30
 Summary: A platform for building and deploying AI agents with structured skills
 Home-page: https://github.com/rosaboyle/airtrain.dev
 Author: Dheeraj Pai
@@ -171,6 +171,15 @@ This project is licensed under the MIT License - see the LICENSE file for detail
 ## Changelog
 
 
+## 0.1.29
+
+- Fixed some issues with th strctured output response example.
+
+## 0.1.28
+
+- Bug fix: reasoning to Fireworks structured output.
+- Added reasoning to Fireworks structured output.
+
 ## 0.1.27
 
 - Added structured completion skills for Fireworks AI

{airtrain-0.1.27 → airtrain-0.1.30}/airtrain/__init__.py

@@ -1,6 +1,6 @@
 """Airtrain - A platform for building and deploying AI agents with structured skills"""
 
-__version__ = "0.1.27"
+__version__ = "0.1.30"
 
 # Core imports
 from .core.skills import Skill, ProcessingError

{airtrain-0.1.27 → airtrain-0.1.30}/airtrain/integrations/fireworks/structured_requests_skills.py

@@ -3,6 +3,7 @@ from pydantic import BaseModel, Field
 import requests
 import json
 from loguru import logger
+import re
 
 from airtrain.core.skills import Skill, ProcessingError
 from airtrain.core.schemas import InputSchema, OutputSchema
@@ -47,6 +48,7 @@ class FireworksStructuredRequestOutput(OutputSchema):
     parsed_response: Any
     used_model: str
     usage: Dict[str, int]
+    reasoning: Optional[str] = None
 
 
 class FireworksStructuredRequestSkill(
@@ -128,19 +130,33 @@ class FireworksStructuredRequestSkill(
                     except json.JSONDecodeError:
                         continue
 
-            # Once complete, parse the full JSON
-            complete_json = "".join(json_buffer)
+            # Once complete, parse the full response with think tags
+            complete_response = "".join(json_buffer)
+            reasoning, json_str = self._parse_response_content(complete_response)
+
             try:
                 parsed_response = input_data.response_model.model_validate_json(
-                    complete_json
+                    json_str
                 )
-                yield {"complete": parsed_response}
+                yield {"complete": parsed_response, "reasoning": reasoning}
             except Exception as e:
                 raise ProcessingError(f"Failed to parse JSON response: {str(e)}")
 
         except Exception as e:
             raise ProcessingError(f"Fireworks streaming request failed: {str(e)}")
 
+    def _parse_response_content(self, content: str) -> tuple[Optional[str], str]:
+        """Parse response content to extract reasoning and JSON."""
+        # Extract reasoning if present
+        reasoning_match = re.search(r"<think>(.*?)</think>", content, re.DOTALL)
+        reasoning = reasoning_match.group(1).strip() if reasoning_match else None
+
+        # Extract JSON
+        json_match = re.search(r"</think>\s*(\{.*\})", content, re.DOTALL)
+        json_str = json_match.group(1).strip() if json_match else content
+
+        return reasoning, json_str
+
     def process(
         self, input_data: FireworksStructuredRequestInput
     ) -> FireworksStructuredRequestOutput:
@@ -150,12 +166,14 @@ class FireworksStructuredRequestSkill(
                 # For streaming, collect and parse the entire response
                 json_buffer = []
                 parsed_response = None
+                reasoning = None
 
                 for chunk in self.process_stream(input_data):
                     if "chunk" in chunk:
                         json_buffer.append(chunk["chunk"])
                     elif "complete" in chunk:
                         parsed_response = chunk["complete"]
+                        reasoning = chunk.get("reasoning")
 
                 if parsed_response is None:
                     raise ProcessingError("Failed to parse streamed response")
@@ -164,6 +182,7 @@ class FireworksStructuredRequestSkill(
                     parsed_response=parsed_response,
                     used_model=input_data.model,
                     usage={},  # Usage stats not available in streaming mode
+                    reasoning=reasoning,
                 )
             else:
                 # For non-streaming, use regular request
@@ -175,14 +194,20 @@ class FireworksStructuredRequestSkill(
                 data = response.json()
 
                 response_content = data["choices"][0]["message"]["content"]
+
+                # Parse the response content to extract reasoning and JSON
+                reasoning, json_str = self._parse_response_content(response_content)
+
+                # Parse the JSON string into the specified model
                 parsed_response = input_data.response_model.model_validate_json(
-                    response_content
+                    json_str
                 )
 
                 return FireworksStructuredRequestOutput(
                     parsed_response=parsed_response,
                     used_model=input_data.model,
                     usage=data["usage"],
+                    reasoning=reasoning,  # Add reasoning to output if present
                 )
 
         except Exception as e:

{airtrain-0.1.27 → airtrain-0.1.30/airtrain.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: airtrain
-Version: 0.1.27
+Version: 0.1.30
 Summary: A platform for building and deploying AI agents with structured skills
 Home-page: https://github.com/rosaboyle/airtrain.dev
 Author: Dheeraj Pai
@@ -171,6 +171,15 @@ This project is licensed under the MIT License - see the LICENSE file for detail
 ## Changelog
 
 
+## 0.1.29
+
+- Fixed some issues with th strctured output response example.
+
+## 0.1.28
+
+- Bug fix: reasoning to Fireworks structured output.
+- Added reasoning to Fireworks structured output.
+
 ## 0.1.27
 
 - Added structured completion skills for Fireworks AI

{airtrain-0.1.27 → airtrain-0.1.30}/changelog.md

@@ -4,6 +4,15 @@ sidebar_position: 9
 
 # Changelog
 
+## 0.1.29
+
+- Fixed some issues with th strctured output response example.
+
+## 0.1.28
+
+- Bug fix: reasoning to Fireworks structured output.
+- Added reasoning to Fireworks structured output.
+
 ## 0.1.27
 
 - Added structured completion skills for Fireworks AI

{airtrain-0.1.27 → airtrain-0.1.30}/examples/integrations/fireworks/structured_requests_example.py

@@ -2,10 +2,11 @@ import sys
 import os
 from pathlib import Path
 from typing import List, Optional
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 from dotenv import load_dotenv
 import time
 import json
+from time import sleep
 
 load_dotenv()
 
@@ -21,13 +22,14 @@ from airtrain.integrations.fireworks.structured_requests_skills import (
 
 
 class MovieReview(BaseModel):
-    """Example schema for movie review"""
+    """Schema for movie review"""
 
-    title: str
-    rating: float
-    pros: List[str]
-    cons: List[str]
-    summary: str
+    title: str = Field(..., description="Movie title")
+    year: int = Field(..., description="Release year")
+    rating: float = Field(..., description="Rating out of 10")
+    review: str = Field(..., description="Detailed review")
+    pros: list[str] = Field(..., description="List of positive points")
+    cons: list[str] = Field(..., description="List of negative points")
 
 
 def stream_example():
@@ -39,21 +41,25 @@ def stream_example():
         system_prompt="You are a movie critic providing structured reviews.",
         model="accounts/fireworks/models/deepseek-r1",
         temperature=0.3,
-        max_tokens=1000,
+        max_tokens=8096,
         response_model=MovieReview,
         stream=True,
     )
 
     print("\nStreaming structured movie review...")
     try:
-        json_buffer = []
         for chunk in skill.process_stream(input_data):
             if "chunk" in chunk:
+                # Print the raw streaming content
                 print(chunk["chunk"], end="", flush=True)
-                time.sleep(0.01)  # Optional: add slight delay for visualization
+                sleep(0.01)  # Optional: add slight delay for visualization
             elif "complete" in chunk:
+                # Print the parsed response and reasoning
                 print("\n\nParsed Response:")
                 print(chunk["complete"].model_dump_json(indent=2))
+                if "reasoning" in chunk:
+                    print("\nReasoning:")
+                    print(chunk["reasoning"])
         print("\n")
     except Exception as e:
         print(f"Error: {str(e)}")
@@ -68,7 +74,7 @@ def non_stream_example():
         system_prompt="You are a movie critic providing structured reviews.",
         model="accounts/fireworks/models/deepseek-r1",
         temperature=0.3,
-        max_tokens=1000,
+        max_tokens=8096,
         response_model=MovieReview,
         stream=False,
     )
@@ -78,6 +84,9 @@ def non_stream_example():
         result = skill.process(input_data)
         print("\nParsed Response:")
         print(result.parsed_response.model_dump_json(indent=2))
+        if result.reasoning:
+            print("\nReasoning:")
+            print(result.reasoning)
         print("\nModel Used:", result.used_model)
         print("Usage Statistics:", result.usage)
     except Exception as e:
@@ -85,10 +94,11 @@ def non_stream_example():
 
 
 def main():
-    # Run streaming example
-    stream_example()
+    """Run the examples"""
+    print("=== Streaming Example ===")
+    # stream_example()
 
-    # Run non-streaming example
+    print("\n=== Non-Streaming Example ===")
     non_stream_example()