Fix: minor bugs and code clean up

flo_ai/flo_ai/arium/builder.py

@@ -341,7 +341,7 @@ def from_yaml(
 
             # Method 4: External file reference
             elif 'yaml_file' in agent_config:
-                agent_builder = AgentBuilder.from_yaml(
+                agent_builder: AgentBuilder = AgentBuilder.from_yaml(
                     yaml_file=agent_config['yaml_file'], base_llm=base_llm
                 )
                 agent = agent_builder.build()
@@ -648,10 +648,8 @@ def _create_agent_from_direct_config(
                         f'Available: {list(available_tools.keys())}'
                     )
 
-        # Extract output schema if present
-        output_schema = agent_config.get('output_schema')
-
         # Handle parser configuration if present
+        output_schema = None
         if 'parser' in agent_config:
             from flo_ai.formatter.yaml_format_parser import FloYamlParser
 
@@ -660,16 +658,17 @@ def _create_agent_from_direct_config(
             parser = FloYamlParser.create(yaml_dict=parser_config)
             output_schema = parser.get_format()
 
-        # Create and return the agent
-        agent = Agent(
-            name=name,
-            system_prompt=job,
-            llm=llm,
-            tools=agent_tools,
-            max_retries=max_retries,
-            reasoning_pattern=reasoning_pattern,
-            output_schema=output_schema,
-            role=role,
+        agent = (
+            AgentBuilder()
+            .with_name(name)
+            .with_prompt(job)
+            .with_llm(llm)
+            .with_tools(agent_tools)
+            .with_retries(max_retries)
+            .with_reasoning(reasoning_pattern)
+            .with_output_schema(output_schema)
+            .with_role(role)
+            .build()
         )
 
         return agent

flo_ai/flo_ai/llm/anthropic_llm.py

@@ -49,7 +49,6 @@ async def generate(
                 'model': self.model,
                 'messages': conversation,
                 'temperature': self.temperature,
-                'max_tokens': 8192,
                 **self.kwargs,
             }
 

flo_ai/flo_ai/llm/gemini_llm.py

@@ -1,9 +1,8 @@
 from typing import Dict, Any, List, Optional
 from google import genai
-import json
+from google.genai import types
 from .base_llm import BaseLLM, ImageMessage
 from flo_ai.tool.base_tool import Tool
-from flo_ai.utils.logger import logger
 
 
 class Gemini(BaseLLM):
@@ -27,9 +26,9 @@ async def generate(
         output_schema: Optional[Dict[str, Any]] = None,
     ) -> Dict[str, Any]:
         # Convert messages to Gemini format
-        # Gemini uses a simple content string format
         contents = []
         system_prompt = ''
+
         for msg in messages:
             role = msg['role']
             message_content = msg['content']
@@ -39,88 +38,65 @@ async def generate(
             else:
                 contents.append(message_content)
 
-        # Add output schema instruction if provided
-        if output_schema:
-            contents += f'\nPlease provide your response in JSON format according to this schema:\n{json.dumps(output_schema, indent=2)}\n'
-
-        # Add function information if provided
-        if functions:
-            contents += f'\nAvailable functions:\n{json.dumps(functions, indent=2)}\n'
-
         try:
             # Prepare generation config
-            generation_config = genai.types.GenerateContentConfig(
+            generation_config = types.GenerateContentConfig(
                 temperature=self.temperature,
                 system_instruction=system_prompt,
                 **self.kwargs,
             )
 
+            # Add tools if functions are provided
+            if functions:
+                tools = types.Tool(function_declarations=functions)
+                generation_config.tools = [tools]
+
+            # Add structured output configuration if output_schema is provided
+            if output_schema:
+                generation_config.response_mime_type = 'application/json'
+                generation_config.response_schema = output_schema
+
             # Make the API call
             response = self.client.models.generate_content(
                 model=self.model,
                 contents=contents,
-                config=generation_config if generation_config else None,
+                config=generation_config,
             )
 
-            # Check if response contains function call information
-            # For now, we'll assume text response and parse for function calls if needed
-            response_text = (
-                response.text if hasattr(response, 'text') else str(response)
-            )
-
-            # Try to detect function calls in the response
-            # This is a simple implementation - in practice, you might need more sophisticated parsing
-            if functions and self._is_function_call_response(response_text):
-                function_call = self._parse_function_call(response_text)
-                if function_call:
+            # Check for function call in the response
+            if (
+                functions
+                and response.candidates
+                and response.candidates[0].content.parts
+            ):
+                part = response.candidates[0].content.parts[0]
+                if hasattr(part, 'function_call') and part.function_call:
+                    function_call = part.function_call
                     return {
-                        'content': response_text,
-                        'function_call': function_call,
+                        'content': response.text,
+                        'function_call': {
+                            'name': function_call.name,
+                            'arguments': function_call.args,
+                        },
                     }
 
+            # Return regular text response
+            response_text = (
+                response.text if hasattr(response, 'text') else str(response)
+            )
             return {'content': response_text}
 
         except Exception as e:
             raise Exception(f'Error in Gemini API call: {str(e)}')
 
-    def _is_function_call_response(self, response_text: str) -> bool:
-        """Check if the response contains a function call"""
-        # Simple heuristic - look for function call patterns
-        return (
-            'function_call' in response_text.lower()
-            or '(' in response_text
-            and ')' in response_text
-        )
-
-    def _parse_function_call(self, response_text: str) -> Optional[Dict[str, Any]]:
-        """Parse function call from response text"""
-        # This is a simplified parser - in practice, you'd want more robust parsing
-        try:
-            # Look for JSON-like function call structure
-            if '{' in response_text and '}' in response_text:
-                # Extract JSON-like content
-                start = response_text.find('{')
-                end = response_text.rfind('}') + 1
-                json_str = response_text[start:end]
-                parsed = json.loads(json_str)
-
-                if 'name' in parsed and 'arguments' in parsed:
-                    return {
-                        'name': parsed['name'],
-                        'arguments': json.dumps(parsed['arguments']),
-                    }
-        except Exception as e:
-            logger.error(f'Error parsing function call: {str(e)}')
-        return None
-
     def get_message_content(self, response: Any) -> str:
         """Extract message content from response"""
         if isinstance(response, dict):
             return response.get('content', '')
         return str(response)
 
     def format_tool_for_llm(self, tool: 'Tool') -> Dict[str, Any]:
-        """Format a single tool for Gemini's API"""
+        """Format a single tool for Gemini's function declarations"""
         return {
             'name': tool.name,
             'description': tool.description,
@@ -142,20 +118,20 @@ def format_tool_for_llm(self, tool: 'Tool') -> Dict[str, Any]:
         }
 
     def format_tools_for_llm(self, tools: List['Tool']) -> List[Dict[str, Any]]:
-        """Format tools for Gemini's API"""
+        """Format tools for Gemini's function declarations"""
         return [self.format_tool_for_llm(tool) for tool in tools]
 
     def format_image_in_message(self, image: ImageMessage) -> str:
         """Format a image in the message"""
         if image.image_file_path:
             with open(image.image_file_path, 'rb') as image_file:
                 image_bytes = image_file.read()
-            return genai.types.Part.from_bytes(
+            return types.Part.from_bytes(
                 data=image_bytes,
                 mime_type=image.mime_type,
             )
         elif image.image_bytes:
-            return genai.types.Part.from_bytes(
+            return types.Part.from_bytes(
                 data=image.image_bytes,
                 mime_type=image.mime_type,
             )

flo_ai/flo_ai/llm/openai_llm.py

@@ -7,7 +7,7 @@
 class OpenAI(BaseLLM):
     def __init__(
         self,
-        model='gpt-40-mini',
+        model='gpt-4o-mini',
         api_key: str = None,
         temperature: float = 0.7,
         base_url: str = None,
@@ -53,6 +53,7 @@ async def generate(
         openai_kwargs = {
             'model': self.model,
             'messages': messages,
+            'temperature': self.temperature,
             **kwargs,
             **self.kwargs,
         }

flo_ai/flo_ai/llm/openai_vllm.py

@@ -18,6 +18,8 @@ def __init__(
             base_url=base_url,
             **kwargs,
         )
+        # Store base_url attribute
+        self.base_url = base_url
 
     # overriden
     async def generate(
@@ -41,6 +43,7 @@ async def generate(
                 )
             else:
                 messages.insert(
+                    0,
                     {
                         'role': 'system',
                         'content': f'Please provide your response in JSON format according to the specified schema.\n \n {output_schema}',

flo_ai/flo_ai/llm/vertexai_llm.py

@@ -17,5 +17,10 @@ def __init__(
     ):
         # Initialize only the BaseLLM part to avoid Gemini's client creation
         BaseLLM.__init__(self, model, api_key, temperature, **kwargs)
+
+        # Store project and location attributes
+        self.project = project
+        self.location = location
+
         # Create VertexAI-specific client
         self.client = genai.Client(project=project, location=location, vertexai=True)

flo_ai/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "flo_ai"
-version = "1.0.0"
+version = "1.0.1-dev1"
 description = "A easy way to create structured AI agents"
 authors = ["rootflo <*@rootflo.ai>"]
 license = "MIT"

flo_ai/setup.py

@@ -5,7 +5,7 @@
 
 setuptools.setup(
     name='flo-ai',
-    version='1.0.0',
+    version='1.0.1-dev1',
     author='Rootflo',
     description='Create composable AI agents',
     long_description=long_description,