airtrain 0.1.27__py3-none-any.whl → 0.1.29__py3-none-any.whl

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.29"
 
 # Core imports
 from .core.skills import Skill, ProcessingError

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.dist-info → airtrain-0.1.29.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: airtrain
-Version: 0.1.27
+Version: 0.1.29
 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,12 @@ This project is licensed under the MIT License - see the LICENSE file for detail
 ## Changelog
 
 
+
+## 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.dist-info → airtrain-0.1.29.dist-info}/RECORD

@@ -1,4 +1,4 @@
-airtrain/__init__.py,sha256=XNINrhMBMu3SCd5wCDb2FLhe3OfbvZqWFEvVwcRKtm0,2099
+airtrain/__init__.py,sha256=KwgaA3OIuQa6v6bg6dmrZp2drOpIlUw6oF08GZjprfA,2099
 airtrain/contrib/__init__.py,sha256=pG-7mJ0pBMqp3Q86mIF9bo1PqoBOVSGlnEK1yY1U1ok,641
 airtrain/contrib/travel/__init__.py,sha256=clmBodw4nkTA-DsgjVGcXfJGPaWxIpCZDtdO-8RzL0M,811
 airtrain/contrib/travel/agents.py,sha256=tpQtZ0WUiXBuhvZtc2JlEam5TuR5l-Tndi14YyImDBM,8975
@@ -25,7 +25,7 @@ airtrain/integrations/fireworks/models.py,sha256=F-MddbLCLAsTjwRr1l6IpJxOegyY4pD
 airtrain/integrations/fireworks/requests_skills.py,sha256=zPIR70l0KdSGmA5WyDEopFAdKHSPltAttJBzWyHu6Bk,5878
 airtrain/integrations/fireworks/skills.py,sha256=OB4epD4CSTxExUCC1oMJ_8rHLOoftlxf0FUoIVrd4mA,5163
 airtrain/integrations/fireworks/structured_completion_skills.py,sha256=IXG4gsZDSfuscrmKIHfnyHkBaCV7zlPInaWXb95iC5k,6428
-airtrain/integrations/fireworks/structured_requests_skills.py,sha256=3JJqyQ3zyWvAhb_cXKGuSKUo7XZduKtQf1814CvY1Co,7244
+airtrain/integrations/fireworks/structured_requests_skills.py,sha256=oRpbKMOcKgY2js16uNkIx866UEwEYSNgFPNYn9cLO3U,8409
 airtrain/integrations/fireworks/structured_skills.py,sha256=BZaLqSOTC11QdZ4kDORS4JnwF_YXBAa-IiwQ5dJiHXw,3895
 airtrain/integrations/google/__init__.py,sha256=ElwgcXfbg_gGMm6zbkMXCQPFKZUb-yTJk986o19A7Cs,214
 airtrain/integrations/google/credentials.py,sha256=KSvWNqW8Mjr4MkysRvUqlrOSGdShNIe5u2OPO6vRrWY,2047
@@ -57,7 +57,7 @@ airtrain/integrations/together/rerank_skill.py,sha256=gjH24hLWCweWKPyyfKZMG3K_g9
 airtrain/integrations/together/schemas.py,sha256=pBMrbX67oxPCr-sg4K8_Xqu1DWbaC4uLCloVSascROg,1210
 airtrain/integrations/together/skills.py,sha256=mUoHc2r5TYQi5iGzwz2aDuUeROGq7teCtNrOlNApef4,6276
 airtrain/integrations/together/vision_models_config.py,sha256=m28HwYDk2Kup_J-a1FtynIa2ZVcbl37kltfoHnK8zxs,1544
-airtrain-0.1.27.dist-info/METADATA,sha256=6TUG7xaOytkupjRdIx_3me_rVaApKtY9T1wZwgvh1jA,5127
-airtrain-0.1.27.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-airtrain-0.1.27.dist-info/top_level.txt,sha256=cFWW1vY6VMCb3AGVdz6jBDpZ65xxBRSqlsPyySxTkxY,9
-airtrain-0.1.27.dist-info/RECORD,,
+airtrain-0.1.29.dist-info/METADATA,sha256=R_xdPE4QQ8KN9dUG4NN-PRiP02GOVKBNSB_Qtemyaig,5243
+airtrain-0.1.29.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+airtrain-0.1.29.dist-info/top_level.txt,sha256=cFWW1vY6VMCb3AGVdz6jBDpZ65xxBRSqlsPyySxTkxY,9
+airtrain-0.1.29.dist-info/RECORD,,