Output parser, and data accumulator for agents, teams & flo

.gitignore

@@ -11,4 +11,5 @@ dataset
 bin
 .DS_Store
 *.sql
-*.log
+*.log
+*.yaml

examples/python/output_parser.py

@@ -0,0 +1,62 @@
+import os
+from flo_ai import FloAgent, FloSession, Flo
+from langchain_community.tools.tavily_search.tool import TavilySearchResults
+from dotenv import load_dotenv
+from langchain_openai import AzureChatOpenAI
+from flo_ai.parsers.flo_json_parser import FloJsonParser
+from flo_ai.state.flo_kv_collector import FloKVCollector
+
+load_dotenv()
+llm = AzureChatOpenAI(
+    azure_endpoint=os.getenv('AZURE_GPT4_ENDPOINT'),
+    model_name='gpt-4o',
+    temperature=0.2,
+    max_tokens=4096,
+    api_version='2024-08-01-preview',
+    api_key=os.getenv('AZURE_OPEN_AI_API_KEY'),
+)
+
+session = FloSession(llm).register_tool(
+    name='TavilySearchResults', tool=TavilySearchResults()
+)
+
+format = {
+    'name': 'NameFormat',
+    'fields': [
+        {
+            'type': 'str',
+            'description': 'The first name of the person',
+            'name': 'first_name',
+        },
+        {
+            'type': 'str',
+            'description': 'The middle name of the person',
+            'name': 'middle_name',
+        },
+        {
+            'type': 'str',
+            'description': 'The last name of the person',
+            'name': 'last_name',
+        },
+    ],
+}
+
+dc = FloKVCollector()
+
+researcher = FloAgent.create(
+    session,
+    name='Researcher',
+    role='Internet Researcher',
+    job='What is the first name, last name  and middle name of the the person user asks about',
+    tools=[TavilySearchResults()],
+    parser=FloJsonParser.create(json_dict=format),
+    data_collector=dc,
+)
+
+
+Flo.set_log_level('DEBUG')
+flo: Flo = Flo.create(session, researcher)
+result = flo.invoke('Mahatma Gandhi')
+
+print(result)
+print(dc.fetch())

examples/python/output_parser_yaml.py

@@ -0,0 +1,52 @@
+import os
+from flo_ai import FloSession, Flo
+from langchain_community.tools.tavily_search.tool import TavilySearchResults
+from dotenv import load_dotenv
+from langchain_openai import AzureChatOpenAI
+from flo_ai.state.flo_kv_collector import FloKVCollector
+
+load_dotenv()
+llm = AzureChatOpenAI(
+    azure_endpoint=os.getenv('AZURE_GPT4_ENDPOINT'),
+    model_name='gpt-4o',
+    temperature=0.2,
+    max_tokens=4096,
+    api_version='2024-08-01-preview',
+    api_key=os.getenv('AZURE_OPEN_AI_API_KEY'),
+)
+
+session = FloSession(llm).register_tool(
+    name='InternetSearchTool', tool=TavilySearchResults()
+)
+
+dc = FloKVCollector()
+
+session.register_data_collector('kv', dc)
+
+simple_reseacher = """
+apiVersion: flo/alpha-v1
+kind: FloAgent
+name: weather-assistant
+agent:
+    name: WeatherAssistant
+    kind: agentic
+    job: >
+      Given the person name, guess the first and last name
+    tools:
+      - name: InternetSearchTool
+    parser:
+        name: NameFormatter
+        fields:
+            - type: str
+              description: The first name of the person
+              name: first_name
+            - type: str
+              description: The first name of the person
+              name: last_name
+    data_collector: kv
+"""
+
+flo: Flo = Flo.build(session, simple_reseacher)
+result = flo.invoke('Gandhi')
+
+print(dc.fetch())

flo_ai/common/flo_logger.py

@@ -66,7 +66,7 @@ def setLevel(self, level: Union[str, int]) -> None:
             handler.setLevel(level)
 
     def _log(
-        self, level: int, msg: str, session: Optional[str] = None, *args, **kwargs
+        self, level: int, msg: str, session: Optional[Any] = None, *args, **kwargs
     ):
         if not self.isEnabledFor(level):
             return

flo_ai/core.py

@@ -18,6 +18,7 @@
     set_logger_internal,
     FloLogConfig,
 )
+from flo_ai.models.flo_node import FloNode
 from flo_ai.models.flo_agent import FloAgent
 from langchain.tools import StructuredTool
 
@@ -84,18 +85,19 @@ def build(
             executable: ExecutableFlo = build_supervised_team(
                 session, to_supervised_team(yaml)
             )
+            if isinstance(executable, FloAgent):
+                executable = FloNode.Builder(session).build_from_agent(executable)
             return Flo(session, executable)
         if routed_team is not None:
             return Flo(session, routed_team.build_routed_team())
         raise FloException("""Either yaml or routed_team should be not None""")
 
     @staticmethod
     def create(session: FloSession, routed_team: Union[FloRouter, FloAgent]):
-        runnable = (
-            routed_team.build_routed_team()
-            if isinstance(routed_team, FloRouter)
-            else routed_team
-        )
+        if isinstance(routed_team, FloRouter):
+            runnable = routed_team.build_routed_team()
+        else:
+            runnable = FloNode.Builder(session).build_from_agent(routed_team)
         return Flo(session, runnable)
 
     @staticmethod

flo_ai/factory/agent_factory.py

@@ -1,6 +1,7 @@
+import json
 from typing import Optional
 from flo_ai.state.flo_session import FloSession
-from flo_ai.yaml.config import AgentConfig
+from flo_ai.yaml.config import AgentConfig, Parser
 from flo_ai.models.flo_agent import FloAgent
 from flo_ai.models.flo_llm_agent import FloLLMAgent
 from flo_ai.models.flo_reflection_agent import FloReflectionAgent
@@ -10,6 +11,7 @@
 from flo_ai.models.delegate import Delegate
 from flo_ai.constants.common_constants import DOCUMENTATION_AGENT_ANCHOR
 from enum import Enum
+from flo_ai.parsers.flo_json_parser import FloJsonParser
 
 
 class AgentKinds(Enum):
@@ -59,7 +61,18 @@ def __create_agentic_agent(
         session: FloSession, agent: AgentConfig, tool_map
     ) -> FloAgent:
         agent_model = AgentFactory.__resolve_model(session, agent.model)
+        dc = (
+            session.data_collectors[agent.data_collector]
+            if agent.data_collector is not None
+            else None
+        )
         tools = [tool_map[tool.name] for tool in agent.tools]
+        if isinstance(agent.parser, Parser):
+            parser = FloJsonParser.create(
+                json_dict=json.loads(agent.parser.model_dump_json())
+            )
+        else:
+            parser = session.parsers[agent.parser] if agent.parser is not None else None
         flo_agent: FloAgent = FloAgent.Builder(
             session,
             name=agent.name,
@@ -69,19 +82,34 @@ def __create_agentic_agent(
             llm=agent_model,
             on_error=session.on_agent_error,
             model_name=agent.model,
+            parser=parser,
+            data_collector=dc,
         ).build()
         return flo_agent
 
     @staticmethod
     def __create_llm_agent(session: FloSession, agent: AgentConfig) -> FloLLMAgent:
         agent_model = AgentFactory.__resolve_model(session, agent.model)
+        dc = (
+            session.data_collectors[agent.data_collector]
+            if agent.data_collector is not None
+            else None
+        )
+        if isinstance(agent.parser, Parser):
+            parser = FloJsonParser.create(
+                json_dict=json.loads(agent.parser.model_dump_json())
+            )
+        else:
+            parser = session.parsers[agent.parser] if agent.parser is not None else None
         builder = FloLLMAgent.Builder(
             session,
             name=agent.name,
             job=agent.job,
             role=agent.role,
             llm=agent_model,
             model_name=agent.model,
+            parser=parser,
+            data_collector=dc,
         )
         llm_agent: FloLLMAgent = builder.build()
         return llm_agent

flo_ai/models/flo_agent.py

@@ -4,10 +4,11 @@
 from langchain_core.runnables import Runnable
 from langchain_core.language_models import BaseLanguageModel
 from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
-from flo_ai.models.flo_executable import ExecutableFlo
+from flo_ai.models.flo_executable import ExecutableFlo, ExecutableType
 from flo_ai.state.flo_session import FloSession
 from typing import Union, Optional, Callable
-from flo_ai.models.flo_executable import ExecutableType
+from flo_ai.state.flo_data_collector import FloDataCollector
+from flo_ai.parsers.flo_parser import FloParser
 
 
 class FloAgent(ExecutableFlo):
@@ -17,11 +18,13 @@ def __init__(
         agent: Runnable,
         executor: AgentExecutor,
         model_name: str,
+        data_collector: Optional[FloDataCollector] = None,
     ) -> None:
         super().__init__(name, executor, ExecutableType.agentic)
         self.model_name = model_name
         self.agent: Runnable = (agent,)
         self.executor: AgentExecutor = executor
+        self.data_collector = data_collector
 
     @staticmethod
     def create(
@@ -32,6 +35,8 @@ def create(
         role: Optional[str] = None,
         on_error: Union[str, Callable] = True,
         llm: Union[BaseLanguageModel, None] = None,
+        parser: Optional[FloParser] = None,
+        data_collector: Optional[FloDataCollector] = None,
     ):
         model_name = 'default' if llm is None else llm.name
         return FloAgent.Builder(
@@ -43,6 +48,8 @@ def create(
             on_error=on_error,
             llm=llm,
             model_name=model_name,
+            parser=parser,
+            data_collector=data_collector,
         ).build()
 
     class Builder:
@@ -57,6 +64,8 @@ def __init__(
             llm: Union[BaseLanguageModel, None] = None,
             on_error: Union[str, Callable] = True,
             model_name: Union[str, None] = 'default',
+            parser: Optional[FloParser] = None,
+            data_collector: Optional[FloDataCollector] = None,
         ) -> None:
             prompt: Union[ChatPromptTemplate, str] = job
             self.name: str = name
@@ -67,16 +76,23 @@ def __init__(
                 if role is not None
                 else [('system', prompt)]
             )
+            if parser is not None:
+                system_prompts.append('\n{format_instructions}')
             system_prompts.append(MessagesPlaceholder(variable_name='messages'))
             system_prompts.append(MessagesPlaceholder(variable_name='agent_scratchpad'))
             self.prompt: ChatPromptTemplate = (
                 ChatPromptTemplate.from_messages(system_prompts)
                 if isinstance(prompt, str)
                 else prompt
             )
+            if parser is not None:
+                self.prompt = self.prompt.partial(
+                    format_instructions=parser.get_format_instructions()
+                )
             self.tools: list[BaseTool] = tools
             self.verbose = verbose
             self.on_error = on_error
+            self.data_collector = data_collector
 
         def build(self) -> AgentExecutor:
             agent = create_tool_calling_agent(self.llm, self.tools, self.prompt)
@@ -87,4 +103,10 @@ def build(self) -> AgentExecutor:
                 return_intermediate_steps=True,
                 handle_parsing_errors=self.on_error,
             )
-            return FloAgent(self.name, agent, executor, model_name=self.model_name)
+            return FloAgent(
+                self.name,
+                agent,
+                executor,
+                model_name=self.model_name,
+                data_collector=self.data_collector,
+            )