flo_ai 0.0.6.dev1__py3-none-any.whl

flo_ai/__init__.py

@@ -0,0 +1,13 @@
+from flo_ai.core import Flo as Flo
+from flo_ai.models.flo_team import FloTeam as FloTeam
+from flo_ai.models.flo_agent import FloAgent as FloAgent
+from flo_ai.router.flo_linear import FloLinear as FloLinear
+from flo_ai.router.flo_router import FloRouter as FloRouter
+from flo_ai.state.flo_session import FloSession as FloSession
+from flo_ai.models.flo_llm_agent import FloLLMAgent as FloLLMAgent
+from flo_ai.models.flo_tool_agent import FloToolAgent as FloToolAgent
+from flo_ai.router.flo_llm_router import FloLLMRouter as FloLLMRouter
+from flo_ai.router.flo_supervisor import FloSupervisor as FloSupervisor
+from flo_ai.retrievers.flo_retriever import FloRagBuilder as FloRagBuilder
+from flo_ai.models.flo_delegation_agent import FloDelegatorAgent as FloDelegatorAgent
+from flo_ai.models.flo_reflection_agent import FloReflectionAgent as FloReflectionAgent

flo_ai/builders/yaml_builder.py

@@ -0,0 +1,65 @@
+from typing import Union
+from flo_ai.models.flo_team import FloTeam
+from flo_ai.yaml.config import (
+    FloRoutedTeamConfig,
+    TeamConfig,
+    AgentConfig,
+    FloAgentConfig,
+)
+from flo_ai.models.flo_executable import ExecutableFlo
+from flo_ai.state.flo_session import FloSession
+from flo_ai.router.flo_router_factory import FloRouterFactory
+from flo_ai.factory.agent_factory import AgentFactory
+from flo_ai.error.flo_exception import FloException
+from flo_ai.yaml.validators import raise_for_name_error
+from flo_ai.common.flo_logger import get_logger
+
+
+def build_supervised_team(
+    session: FloSession, flo_config: Union[FloRoutedTeamConfig, FloAgentConfig]
+) -> ExecutableFlo:
+    name_set = set()
+    if isinstance(flo_config, FloRoutedTeamConfig):
+        team_config: TeamConfig = flo_config.team
+        team = parse_and_build_subteams(session, team_config, name_set)
+        return team
+    elif isinstance(flo_config, FloAgentConfig):
+        agent_config: AgentConfig = flo_config.agent
+        validate_names(name_set, agent_config.name, session)
+        agent = AgentFactory.create(session, agent_config)
+        return agent
+
+
+def validate_team(name_set: set, team_config: TeamConfig, session: FloSession):
+    validate_names(name_set, team_config.name, session)
+
+
+def parse_and_build_subteams(
+    session: FloSession, team_config: TeamConfig, name_set=set()
+) -> ExecutableFlo:
+    flo_team = None
+    validate_team(name_set, team_config, session)
+    if team_config.agents:
+        members = [AgentFactory.create(session, agent) for agent in team_config.agents]
+        flo_team = FloTeam.Builder(session, team_config.name, members=members).build()
+        router = FloRouterFactory.create(session, team_config, flo_team)
+        flo_routed_team = router.build_routed_team()
+    else:
+        flo_teams = []
+        for subteam in team_config.subteams:
+            flo_subteam = parse_and_build_subteams(session, subteam, name_set)
+            flo_teams.append(flo_subteam)
+        flo_team = FloTeam.Builder(session, team_config.name, members=flo_teams).build()
+        router = FloRouterFactory.create(session, team_config, flo_team)
+        flo_routed_team = router.build_routed_team()
+    return flo_routed_team
+
+
+def validate_names(name_set: set, name, session: FloSession):
+    raise_for_name_error(name)
+    if name in name_set:
+        get_logger().error(f"Duplicate name found: '{name}'", session)
+        raise FloException(
+            f"The name '{name}' is duplicate in the config. Make sure all teams and agents have unique names"
+        )
+    name_set.add(name)

flo_ai/callbacks/__init__.py

@@ -0,0 +1,15 @@
+from flo_ai.callbacks.flo_callbacks import (
+    flo_agent_callback,
+    flo_router_callback,
+    flo_tool_callback,
+    flo_call_back,
+)
+from flo_ai.callbacks.flo_execution_logger import FloExecutionLogger
+
+__all__ = [
+    'flo_agent_callback',
+    'flo_router_callback',
+    'flo_tool_callback',
+    'flo_call_back',
+    'FloExecutionLogger',
+]

flo_ai/callbacks/flo_callbacks.py

@@ -0,0 +1,223 @@
+from typing import Any, Union, Callable, Optional, Dict
+from dataclasses import dataclass, field
+from flo_ai.common.flo_logger import get_logger
+
+
+@dataclass
+class FloCallbackResponse:
+    type: str
+    name: Optional[str] = None
+    model_name: Optional[str] = None
+    input: Optional[str] = None
+    output: Optional[str] = None
+    error: Union[Exception, KeyboardInterrupt, None] = None
+    args: Dict = field(default_factory=dict)
+
+
+class FloToolCallback:
+    def __init__(self) -> None:
+        pass
+
+    def on_tool_start(
+        self, name: str, input: Any, **kwargs: Any
+    ) -> Optional[FloCallbackResponse]:
+        pass
+
+    def on_tool_end(
+        self, name: str, output: Any, **kwargs: Any
+    ) -> Optional[FloCallbackResponse]:
+        pass
+
+    def on_tool_error(
+        self, name: str, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
+    ) -> Optional[FloCallbackResponse]:
+        pass
+
+
+class FloAgentCallback:
+    def __init__(self) -> None:
+        pass
+
+    def on_agent_start(
+        self, name: str, model_name: str, input: Any, **kwargs: Any
+    ) -> Optional[FloCallbackResponse]:
+        pass
+
+    def on_agent_end(
+        self, name: str, model_name: str, output: Any, **kwargs: Any
+    ) -> Optional[FloCallbackResponse]:
+        pass
+
+    def on_agent_error(
+        self,
+        name: str,
+        model_name: str,
+        error: Union[Exception, KeyboardInterrupt],
+        **kwargs: Any,
+    ) -> Optional[FloCallbackResponse]:
+        pass
+
+
+class FloRouterCallback:
+    def __init__(self) -> None:
+        pass
+
+    def on_router_start(
+        self, name: str, model_name: str, input: Any, **kwargs: Any
+    ) -> Optional[FloCallbackResponse]:
+        pass
+
+    def on_router_end(
+        self, name: str, model_name: str, output: Any, **kwargs: Any
+    ) -> Optional[FloCallbackResponse]:
+        pass
+
+    def on_router_error(
+        self,
+        name: str,
+        model_name: str,
+        error: Union[Exception, KeyboardInterrupt],
+        **kwargs: Any,
+    ) -> None:
+        Optional[FloCallbackResponse]
+
+
+def safe_call_cb(func, cb_response: FloCallbackResponse, ignore_error=True):
+    try:
+        func(cb_response)
+    except Exception as e:
+        if ignore_error:
+            get_logger().warning(e)
+        else:
+            raise e
+
+
+class FunctionalFloToolCallbackImpl(FloToolCallback):
+    def __init__(self, func: Callable, ignore_error: bool = True) -> None:
+        super().__init__()
+        self.func = func
+        self.ignore_error = ignore_error
+
+    def on_tool_start(
+        self, name: str, input: Any, **kwargs: Any
+    ) -> Optional[FloCallbackResponse]:
+        cb_response = FloRouterCallback('on_tool_start', name, input=input, args=kwargs)
+        safe_call_cb(self.func, cb_response, self.ignore_error)
+        return cb_response
+
+    def on_tool_end(
+        self, name: str, output: Any, **kwargs: Any
+    ) -> Optional[FloCallbackResponse]:
+        cb_response = FloRouterCallback('on_tool_end', name, output=output, args=kwargs)
+        safe_call_cb(self.func, cb_response, self.ignore_error)
+        return cb_response
+
+    def on_tool_error(
+        self, name: str, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
+    ) -> Optional[FloCallbackResponse]:
+        cb_response = FloRouterCallback('on_tool_error', name, error=error, args=kwargs)
+        safe_call_cb(self.func, cb_response, self.ignore_error)
+        return cb_response
+
+
+class FunctionalFloAgentCallbackImpl(FloAgentCallback):
+    def __init__(self, func: Callable, ignore_error: bool = True) -> None:
+        super().__init__()
+        self.func = func
+        self.ignore_error = ignore_error
+
+    def on_agent_start(
+        self, name: str, model_name: str, input: Any, **kwargs: Any
+    ) -> Any:
+        cb_response = FloCallbackResponse(
+            'on_agent_start', name, input=input, args=kwargs, model_name=model_name
+        )
+        safe_call_cb(self.func, cb_response, self.ignore_error)
+        return cb_response
+
+    def on_agent_end(
+        self, name: str, model_name: str, output: Any, **kwargs: Any
+    ) -> None:
+        cb_response = FloCallbackResponse(
+            'on_agent_end', name, output=output, args=kwargs, model_name=model_name
+        )
+        safe_call_cb(self.func, cb_response, self.ignore_error)
+        return cb_response
+
+    def on_agent_error(
+        self,
+        name: str,
+        model_name: str,
+        error: Union[Exception, KeyboardInterrupt],
+        **kwargs: Any,
+    ) -> None:
+        cb_response = FloCallbackResponse(
+            'on_agent_error', name, error=error, args=kwargs, model_name=model_name
+        )
+        safe_call_cb(self.func, cb_response, self.ignore_error)
+        return cb_response
+
+
+class FunctionalFloRouterCallbackImpl(FloRouterCallback):
+    def __init__(self, func: Callable, ignore_error: bool = True) -> None:
+        super().__init__()
+        self.func = func
+        self.ignore_error = ignore_error
+
+    def on_router_start(
+        self, name: str, model_name: str, input: Any, **kwargs: Any
+    ) -> Any:
+        cb_response = FloCallbackResponse(
+            'on_router_start', name, input=input, args=kwargs, model_name=model_name
+        )
+        safe_call_cb(self.func, cb_response, self.ignore_error)
+        return cb_response
+
+    def on_router_end(
+        self, name: str, model_name: str, output: Any, **kwargs: Any
+    ) -> None:
+        cb_response = FloCallbackResponse(
+            'on_router_end', name, output=output, args=kwargs, model_name=model_name
+        )
+        safe_call_cb(self.func, cb_response, self.ignore_error)
+        return cb_response
+
+    def on_router_error(
+        self,
+        name: str,
+        model_name: str,
+        error: Union[Exception, KeyboardInterrupt],
+        **kwargs: Any,
+    ) -> None:
+        cb_response = FloCallbackResponse(
+            'on_router_error', name, error=error, args=kwargs, model_name=model_name
+        )
+        safe_call_cb(self.func, cb_response, self.ignore_error)
+        return cb_response
+
+
+class FloCallback(
+    FunctionalFloToolCallbackImpl,
+    FunctionalFloAgentCallbackImpl,
+    FunctionalFloRouterCallbackImpl,
+):
+    def __init__(self, func: Callable, ignore_error: bool = True) -> None:
+        FunctionalFloToolCallbackImpl.__init__(self, func, ignore_error)
+        FunctionalFloAgentCallbackImpl.__init__(self, func, ignore_error)
+        FunctionalFloRouterCallbackImpl.__init__(self, func, ignore_error)
+
+
+def flo_tool_callback(func: Callable, ignore_error=True) -> FloToolCallback:
+    return FunctionalFloToolCallbackImpl(func, ignore_error)
+
+
+def flo_agent_callback(func: Callable, ignore_error=True) -> FloToolCallback:
+    return FunctionalFloAgentCallbackImpl(func, ignore_error)
+
+
+def flo_router_callback(func: Callable, ignore_error=True) -> FloRouterCallback:
+    return FunctionalFloRouterCallbackImpl(func, ignore_error)
+
+
+def flo_call_back(func: Callable, ignore_error=True) -> FloRouterCallback:
+    return FloCallback(func, ignore_error)

flo_ai/callbacks/flo_execution_logger.py

@@ -0,0 +1,277 @@
+import json
+from typing import Any, Dict, Optional
+from datetime import datetime
+from uuid import UUID
+from langchain_core.callbacks import BaseCallbackHandler
+from langchain.schema.agent import AgentAction, AgentFinish
+from langchain.schema import HumanMessage, AIMessage, BaseMessage
+from langchain_core.prompts.chat import ChatPromptValue
+from flo_ai.storage.data_collector import DataCollector
+from flo_ai.common.flo_logger import get_logger
+from abc import ABC, abstractmethod
+
+
+class ToolLogger(ABC):
+    @abstractmethod
+    def log_all_tools(session_tools):
+        pass
+
+
+class EnhancedJSONEncoder(json.JSONEncoder):
+    def default(self, obj):
+        if isinstance(obj, (HumanMessage, AIMessage, BaseMessage)):
+            return {
+                'type': obj.__class__.__name__,
+                'content': obj.content,
+                'additional_kwargs': obj.additional_kwargs,
+            }
+        elif isinstance(obj, AgentAction):
+            return {
+                'type': 'AgentAction',
+                'tool': obj.tool,
+                'tool_input': obj.tool_input,
+                'log': obj.log,
+            }
+        elif isinstance(obj, AgentFinish):
+            return {
+                'type': 'AgentFinish',
+                'return_values': obj.return_values,
+                'log': obj.log,
+            }
+        elif isinstance(obj, ChatPromptValue):
+            return {
+                'type': 'ChatPromptValue',
+                'messages': [self.default(msg) for msg in obj.messages],
+            }
+        elif isinstance(obj, datetime):
+            return obj.isoformat()
+        elif isinstance(obj, UUID):
+            return str(obj)
+        elif hasattr(obj, 'to_dict'):
+            return obj.to_dict()
+        return super().default(obj)
+
+
+class FloExecutionLogger(BaseCallbackHandler, ToolLogger):
+    def __init__(self, data_collector: DataCollector):
+        self.data_collector = data_collector
+        self.runs = {}
+        self.encoder = EnhancedJSONEncoder()
+        self.query = None
+        self.added_tools = set()
+        self.prompt = {}
+
+    def _encode_entry(self, entry: Dict[str, Any]) -> Dict[str, Any]:
+        return json.loads(self.encoder.encode(entry))
+
+    def _store_entry(self, entry: Dict[str, Any]) -> None:
+        try:
+            encoded_entry = self._encode_entry(entry)
+            self.data_collector.store_log(encoded_entry)
+        except Exception as e:
+            get_logger().error(f'Error storing entry in FloExecutionLogger: {e}')
+
+    def on_llm_start(
+        self,
+        serialized: dict[str, Any],
+        prompts: list[str],
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        tags: Optional[list[str]] = None,
+        metadata: Optional[dict[str, Any]] = None,
+        **kwargs: Any,
+    ) -> None:
+        self.prompt[str(run_id)] = prompts
+
+    def on_chain_start(
+        self,
+        serialized: Dict[str, Any],
+        inputs: Dict[str, Any],
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        **kwargs: Any,
+    ) -> None:
+        chain_name = (
+            serialized.get('name', 'unnamed_chain') if serialized else 'unnamed_chain'
+        )
+
+        if parent_run_id and chain_name != 'agent_chain':
+            return
+        if isinstance(inputs, dict):
+            user_input = inputs.get('messages', {})
+        else:
+            user_input = {}
+        if (
+            user_input
+            and len(user_input) > 0
+            and isinstance(user_input[0], HumanMessage)
+        ):
+            if isinstance(user_input[0], HumanMessage):
+                self.query = user_input[0].content
+
+        self.runs[str(run_id)] = {
+            'type': 'chain',
+            'start_time': datetime.utcnow(),
+            'inputs': inputs,
+            'name': chain_name,
+            'run_id': str(run_id),
+            'parent_run_id': str(parent_run_id) if parent_run_id else None,
+        }
+
+    def on_chain_end(
+        self,
+        outputs: Dict[str, Any],
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        **kwargs: Any,
+    ) -> None:
+        if str(run_id) in self.runs:
+            run_info = self.runs[str(run_id)]
+            if run_info['type'] != 'chain' and run_info['type'] != 'llm':
+                return
+            run_info['end_time'] = datetime.utcnow()
+            run_info['outputs'] = outputs
+            run_info['status'] = 'completed'
+            run_info['parent_run_id'] = str(parent_run_id) if parent_run_id else None
+            run_info['prompt'] = (
+                self.prompt[str(run_id)] if str(run_id) in self.prompt else []
+            )
+            self._store_entry(run_info)
+            del self.runs[str(run_id)]
+        else:
+            if isinstance(outputs, ChatPromptValue) or isinstance(outputs, AgentFinish):
+                run_info = {}
+                run_info['type'] = 'llm'
+                run_info['end_time'] = datetime.utcnow()
+                run_info['inputs'] = outputs
+                run_info['status'] = 'completed'
+                run_info['run_id'] = str(run_id)
+                run_info['parent_run_id'] = (
+                    str(parent_run_id) if parent_run_id else None
+                )
+                self.runs[str(parent_run_id)] = run_info
+
+    def on_chain_error(
+        self,
+        error: Exception,
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        **kwargs: Any,
+    ) -> None:
+        pass
+
+    def on_tool_start(
+        self,
+        serialized: Dict[str, Any],
+        input_str: str,
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        tags: Optional[list[str]] = None,
+        metadata: Optional[dict[str, Any]] = None,
+        inputs: Optional[dict[str, Any]] = None,
+        **kwargs: Any,
+    ) -> None:
+        self.runs[str(run_id)] = {
+            'type': 'tool',
+            'query': self.query,
+            'start_time': datetime.utcnow(),
+            'tool_name': serialized.get('name', 'unnamed_tool'),
+            'input': input_str,
+            'parent_run_id': str(parent_run_id) if parent_run_id else None,
+        }
+
+    def on_tool_end(
+        self,
+        output: str,
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        tags: Optional[list[str]] = None,
+        **kwargs: Any,
+    ) -> None:
+        if str(run_id) in self.runs:
+            run_info = self.runs[str(run_id)]
+            run_info['end_time'] = datetime.utcnow()
+            run_info['output'] = output
+            run_info['status'] = 'completed'
+            self._store_entry(run_info)
+            del self.runs[str(run_id)]
+
+    def on_tool_error(
+        self,
+        error: Exception,
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        **kwargs: Any,
+    ) -> None:
+        if str(run_id) in self.runs:
+            run_info = self.runs[str(run_id)]
+            run_info['end_time'] = datetime.utcnow()
+            run_info['error'] = str(error)
+            run_info['status'] = 'error'
+            self._store_entry(run_info)
+            del self.runs[str(run_id)]
+
+    def on_agent_action(
+        self,
+        action: AgentAction,
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        **kwargs: Any,
+    ) -> None:
+        agent_info = {
+            'type': 'agent_action',
+            'start_time': datetime.utcnow(),
+            'tool': action.tool,
+            'tool_input': action.tool_input,
+            'log': action.log,
+            'parent_run_id': str(parent_run_id) if parent_run_id else None,
+        }
+        self.runs[str(run_id)] = agent_info
+        self._store_entry(agent_info)
+
+    def on_agent_finish(
+        self,
+        finish: AgentFinish,
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        **kwargs: Any,
+    ) -> None:
+        log_entry = {
+            'type': 'agent_finish',
+            'time': datetime.utcnow(),
+            'output': finish.return_values,
+            'log': finish.log,
+            'parent_run_id': str(parent_run_id) if parent_run_id else None,
+        }
+        self._store_entry(log_entry)
+
+    def log_all_tools(self, session_tools):
+        try:
+            tools = []
+
+            for val in session_tools:
+                tool_name = session_tools[val].name
+                if tool_name not in self.added_tools:
+                    tools.append(
+                        {
+                            'tool_name': tool_name,
+                            'description': session_tools.get(val).description,
+                            'args': session_tools.get(val).args,
+                        }
+                    )
+                    self.added_tools.add(tool_name)
+
+            encoded_entry = self._encode_entry(tools)
+            if encoded_entry:
+                self.data_collector.store_tool_log(encoded_entry)
+        except Exception as e:
+            get_logger().error(f'Error storing tool in FloExecutionLogger: {e}')

flo_ai/common/flo_langchain_logger.py

@@ -0,0 +1,68 @@
+from typing import Any, Dict, List, Union
+from langchain.callbacks.base import BaseCallbackHandler
+from langchain.schema import AgentAction, AgentFinish, LLMResult
+from flo_ai.common.flo_logger import get_logger
+from flo_ai.callbacks.flo_callbacks import FloToolCallback
+
+
+class FloLangchainLogger(BaseCallbackHandler):
+    def __init__(self, session_id: str, tool_callbacks: List[FloToolCallback] = []):
+        self.session_id = session_id
+        self.tool_callbacks = tool_callbacks
+
+    def on_llm_start(
+        self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
+    ) -> None:
+        get_logger().debug(f'onLLMStart: {prompts}', self)
+
+    def on_llm_new_token(self, token: str, **kwargs: Any) -> None:
+        get_logger().debug(f'onNewToken: {token}', self)
+
+    def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
+        get_logger().debug(f'onLLMEnd: {response.generations}', self)
+
+    def on_llm_error(
+        self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
+    ) -> None:
+        get_logger().debug(f'onLLMEnd: {error}', self)
+
+    def on_chain_start(
+        self, serialized: Dict[str, Any], inputs: Dict[str, Any], **kwargs: Any
+    ) -> None:
+        get_logger().debug(f'onChainStart: {inputs}', self)
+
+    def on_chain_end(self, outputs: Dict[str, Any], **kwargs: Any) -> None:
+        get_logger().debug(f'onChainEnd: {outputs}', self)
+
+    def on_chain_error(
+        self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
+    ) -> None:
+        get_logger().debug(f'onChainError: {error}', self)
+
+    def on_tool_start(
+        self, serialized: Dict[str, Any], input_str: str, **kwargs: Any
+    ) -> None:
+        get_logger().debug(f'onToolStart: {input_str}', self)
+        [
+            x.on_tool_start(serialized['name'], kwargs['inputs'], kwargs)
+            for x in self.tool_callbacks
+        ]
+
+    def on_tool_end(self, output: str, **kwargs: Any) -> None:
+        get_logger().debug(f'onToolEnd: {output}', self)
+        [x.on_tool_end(kwargs['name'], output, kwargs) for x in self.tool_callbacks]
+
+    def on_tool_error(
+        self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
+    ) -> None:
+        get_logger().debug(f'onToolError: {error}', self)
+        [x.on_tool_error(kwargs['name'], error, kwargs) for x in self.tool_callbacks]
+
+    def on_text(self, text: str, **kwargs: Any) -> None:
+        get_logger().debug(f'onText: {text}', self)
+
+    def on_agent_action(self, action: AgentAction, **kwargs: Any) -> Any:
+        get_logger().debug(f'onAgentAction: {action.tool} - {action.tool_input}', self)
+
+    def on_agent_finish(self, finish: AgentFinish, **kwargs: Any) -> None:
+        get_logger().debug(f'onAgentFinish: {finish.return_values}', self)