flo-ai 0.0.5.dev3__tar.gz → 0.0.5.dev5__tar.gz

{flo_ai-0.0.5.dev3 → flo_ai-0.0.5.dev5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: flo-ai
-Version: 0.0.5.dev3
+Version: 0.0.5.dev5
 Summary: A easy way to create structured AI agents
 License: MIT
 Author: vizsatiz

{flo_ai-0.0.5.dev3 → flo_ai-0.0.5.dev5}/flo_ai/models/flo_node.py

@@ -25,15 +25,22 @@ class FloNode:
         name: str,
         kind: ExecutableType,
         delegate: Optional[Delegate] = None,
+        async_func: functools.partial = None,
     ) -> None:
         self.name = name
         self.func = func
         self.kind: ExecutableType = kind
         self.delegate = delegate
+        self.async_func = async_func
 
     def invoke(self, query, config):
         return self.func({STATE_NAME_MESSAGES: [HumanMessage(content=query)]})
 
+    async def ainvoke(self, query, config):
+        return await self.async_func(
+            {STATE_NAME_MESSAGES: [HumanMessage(content=query)]}
+        )
+
     class Builder:
         def __init__(self, session: FloSession) -> None:
             self.session = session
@@ -47,7 +54,17 @@ class FloNode:
                 model_name=flo_agent.model_name,
                 data_collector=flo_agent.data_collector,
             )
-            return FloNode(agent_func, flo_agent.name, flo_agent.type)
+            agent_func_async = functools.partial(
+                FloNode.Builder.__async_teamflo_agent_node,
+                agent=flo_agent.runnable,
+                name=flo_agent.name,
+                session=self.session,
+                model_name=flo_agent.model_name,
+                data_collector=flo_agent.data_collector,
+            )
+            return FloNode(
+                agent_func, flo_agent.name, flo_agent.type, async_func=agent_func_async
+            )
 
         def build_from_reflection(self, flo_agent: FloReflectionAgent) -> 'FloNode':
             agent_func = functools.partial(
@@ -154,6 +171,57 @@ class FloNode:
             ]
             return {STATE_NAME_MESSAGES: [AIMessage(content=output, name=name)]}
 
+        @staticmethod
+        async def __async_teamflo_agent_node(
+            state: TeamFloAgentState,
+            agent: AgentExecutor,
+            name: str,
+            session: FloSession,
+            model_name: str,
+            data_collector: Optional[FloDataCollector] = None,
+        ):
+            agent_cbs: List[FloAgentCallback] = FloNode.Builder.__filter_callbacks(
+                session, FloAgentCallback
+            )
+            flo_cbs: List[FloCallback] = FloNode.Builder.__filter_callbacks(
+                session, FloCallback
+            )
+            [
+                callback.on_agent_start(name, model_name, state['messages'], **{})
+                for callback in agent_cbs
+            ]
+            [
+                callback.on_agent_start(name, model_name, state['messages'], **{})
+                for callback in flo_cbs
+            ]
+            try:
+                result = await agent.ainvoke(state)
+                output = result if isinstance(result, str) else result['output']
+                if data_collector is not None:
+                    get_logger().info(
+                        'appending output to data collector', session=session
+                    )
+                    data_collector.append(output)
+            except Exception as e:
+                [
+                    callback.on_agent_error(name, model_name, e, **{})
+                    for callback in agent_cbs
+                ]
+                [
+                    callback.on_agent_error(name, model_name, e, **{})
+                    for callback in flo_cbs
+                ]
+                raise e
+            [
+                callback.on_agent_end(name, model_name, output, **{})
+                for callback in agent_cbs
+            ]
+            [
+                callback.on_agent_start(name, model_name, output, **{})
+                for callback in flo_cbs
+            ]
+            return {STATE_NAME_MESSAGES: [AIMessage(content=output, name=name)]}
+
         @staticmethod
         def __filter_callbacks(session: FloSession, type: Type):
             cbs = session.callbacks

{flo_ai-0.0.5.dev3 → flo_ai-0.0.5.dev5}/flo_ai/parsers/flo_json_parser.py

@@ -18,6 +18,7 @@ class ParseContract:
 class FloJsonParser(FloParser):
     def __init__(self, parse_contract: ParseContract):
         self.contract = parse_contract
+        self._cached_models = {}
         super().__init__()
 
     def __dict_list_to_csv_string(self, data):
@@ -35,30 +36,76 @@ class FloJsonParser(FloParser):
 
         return f'```\n{csv_string}```'
 
-    def __create_contract_from_json(self) -> BaseModel:
+    def __create_nested_model(
+        self, field_def: Dict[str, Any], model_name: str
+    ) -> BaseModel:
+        """Creates a nested Pydantic model for object types"""
+        if model_name in self._cached_models:
+            return self._cached_models[model_name]
+
+        nested_fields = {}
+        for nested_field in field_def['fields']:
+            nested_type = self.__get_field_type_annotation(
+                nested_field, f"{model_name}_{nested_field['name']}"
+            )
+            field_description = nested_field['description']
+            nested_fields[nested_field['name']] = (
+                nested_type,
+                Field(..., description=field_description),
+            )
+
+        NestedModel = create_model(model_name, **nested_fields)
+        self._cached_models[model_name] = NestedModel
+        return NestedModel
+
+    def __get_field_type_annotation(
+        self, field: Dict[str, Any], model_name: str
+    ) -> Any:
+        """Determines the type annotation for a field, handling nested objects"""
         type_mapping = {
             'str': str,
             'int': int,
             'bool': bool,
             'float': float,
-            'literal': Literal,
+            'literal': self.__create_literal_type,
+            'object': lambda f: self.__create_nested_model(f, model_name),
+            'array': lambda f: List[
+                self.__get_field_type_annotation(f['items'], f'{model_name}_item')
+            ],
         }
+
+        field_type = field['type']
+        type_handler = type_mapping.get(field_type)
+
+        if type_handler is None:
+            raise ValueError(f'Unsupported type: {field_type}')
+
+        return (
+            type_handler(field)
+            if field_type in ['literal', 'object', 'array']
+            else type_handler
+        )
+
+    def __create_literal_type(self, field: Dict[str, Any]) -> Any:
+        """Creates a Literal type from field definition"""
+        literal_values = field.get('values', [])
+        if not literal_values:
+            raise ValueError(
+                f"Field '{field['name']}' of type 'literal' must specify 'values'."
+            )
+        literals = [literal_value['value'] for literal_value in literal_values]
+        return Literal[tuple(literals)]
+
+    def __create_contract_from_json(self) -> BaseModel:
         pydantic_fields = {}
         for field in self.contract.fields:
-            field_type = field['type']
-            if field_type == 'literal':
+            field_type = self.__get_field_type_annotation(
+                field, f"{self.contract.name}_{field['name']}"
+            )
+
+            if field['type'] == 'literal':
                 literal_values = field.get('values', [])
-                if not literal_values:
-                    raise ValueError(
-                        f"Field '{field['name']}' of type 'literal' must specify 'values'."
-                    )
-                literals = [literal_value['value'] for literal_value in literal_values]
-                field_type_annotation = Literal[tuple(literals)]
-                default_prompt = (
-                    field['default_value_prompt']
-                    if 'default_value_prompt' in field
-                    else ''
-                )
+                default_prompt = field.get('default_value_prompt', '')
                 field_description = f"""
                 {field['description']}
                 Following are the list of possibles values and its correponding description:
@@ -68,13 +115,10 @@ class FloJsonParser(FloParser):
                 {default_prompt}
                 """
             else:
-                field_type_annotation = type_mapping.get(field_type)
-                if field_type_annotation is None:
-                    raise ValueError(f'Unsupported type: {field_type}')
                 field_description = field['description']
 
             pydantic_fields[field['name']] = (
-                field_type_annotation,
+                field_type,
                 Field(..., description=field_description),
             )
 
@@ -86,6 +130,7 @@ class FloJsonParser(FloParser):
             pydantic_object=self.__create_contract_from_json()
         ).get_format_instructions()
 
+    @staticmethod
     def create(json_dict: Optional[Dict] = None, json_path: Optional[str] = None):
         return FloJsonParser.Builder(json_dict=json_dict, json_path=json_path).build()
 

{flo_ai-0.0.5.dev3 → flo_ai-0.0.5.dev5}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "flo-ai"
-version = "0.0.5-dev3"
+version = "0.0.5-dev5"
 description = "A easy way to create structured AI agents"
 authors = ["vizsatiz <vishnu@rootfor.xyz>"]
 license = "MIT"