Subtle fixes after extending support to python 3.9

README.md

@@ -233,6 +233,114 @@ session.register_tool(name='Adder', tool=addition_tool)
 
 **Note:** `@flotool` comes with inherent error handling capabilities to retry if an exception is thrown. Use `unsafe=True` to disable error handling
 
+## Output Parsing and formatting
+
+FloAI now supports output parsing using JSON or YAML formatter. You can now defined your output formatter using `pydantic` and use the same in code or directly make it part of the Agent Definition Yaml (ADY)
+
+### Using Agent Defintion YAML
+
+We have added parser key to your agent schema, which gives you the output. The following is the schema of the parser
+
+```yaml
+name: SchemaName
+fields:
+  - name: field_name
+    type: data_type
+    description: field_description
+    values: <optional(for literals, all possible values that can be taken)>
+      - value: <the value>
+        description: value_description
+```
+
+### Supported Field Types
+
+#### Primitive Types
+
+- str: String values
+- int: Integer values
+- bool: Boolean values
+- float: Floating-point values
+
+##### Complex Types
+
+- array: Lists of items
+- object: Nested objects
+- literal: Enumerated values
+
+
+Here an example of a simple summarization agent yaml that produces output a structured manner.
+
+```yaml
+apiVersion: flo/alpha-v1
+kind: FloAgent
+name: SummarizationFlo
+agent:
+  name: SummaryAgent
+  kind: llm
+  role: Book summarizer agent
+  job: >
+    You are an given a paragraph from a book
+    and your job is to understand the information in it and extract summary
+  parser:
+    name: BookSummary
+    fields:
+      - name: long_summary
+        type: str
+        description: A comprehensive summary of the book, with all the major topics discussed
+      - name: short_summary
+        type: str
+        description: A short summary of the book in less than 20 words
+```
+
+As you can see here, the `parser` key makes sure that output of this agent will be the given key value format.
+
+### Using parser with code
+
+You can define parser as json in code and use it easily, here is an example:
+
+```python
+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': 'literal',
+            'description': 'The last name of the person, the value can be either of Vishnu or Satis',
+            'name': 'last_name',
+            'values': [
+                {'value': 'Vishnu', 'description': 'If the first_name starts with K'},
+                {'value': 'Satis', 'description': 'If the first_name starts with M'},
+            ],
+            'default_value_prompt': 'If none of the above value is suited, please use value other than the above in snake-case',
+        },
+    ],
+}
+
+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)
+)
+
+
+Flo.set_log_level('DEBUG')
+flo: Flo = Flo.create(session, researcher)
+result = flo.invoke('Mahatma Gandhi')
+
+```
+
 ## 📊 Tool Logging and Data Collection
 
 FloAI provides built-in capabilities for logging tool calls and collecting data through the `FloExecutionLogger` and `DataCollector` classes, facilitating the creation of valuable training data.

examples/agentic_rag.ipynb

@@ -243,7 +243,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.6"
+   "version": "3.9.19"
   }
  },
  "nbformat": 4,

examples/email_reply_agent.ipynb

@@ -415,7 +415,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.6"
+   "version": "3.9.19"
   }
  },
  "nbformat": 4,

examples/python/output_parser.py

@@ -3,8 +3,8 @@
 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
+from flo_ai.parsers import FloJsonParser
+from flo_ai.state import FloJsonOutputCollector
 
 load_dotenv()
 llm = AzureChatOpenAI(
@@ -46,7 +46,7 @@
     ],
 }
 
-dc = FloKVCollector()
+dc = FloJsonOutputCollector()
 
 researcher = FloAgent.create(
     session,

examples/python/output_parser_yaml.py

@@ -3,7 +3,7 @@
 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
+from flo_ai.state import FloJsonOutputCollector
 
 load_dotenv()
 llm = AzureChatOpenAI(
@@ -19,9 +19,9 @@
     name='InternetSearchTool', tool=TavilySearchResults()
 )
 
-dc = FloKVCollector()
+dc = FloJsonOutputCollector()
 
-session.register_data_collector('kv', dc)
+session.register_output_collector('kv', dc)
 
 simple_reseacher = """
 apiVersion: flo/alpha-v1

examples/python/tool_agent.py

@@ -14,8 +14,8 @@
 
 
 class PrintStateTool(BaseTool):
-    name = 'printStateTool'
-    description = 'Just print the state'
+    name: str = 'printStateTool'
+    description: str = 'Just print the state'
 
     def _run(self, **kwargs) -> str:
         return 'Print tool call success'

flo_ai/core.py

@@ -127,7 +127,8 @@ def set_logger(logging_config: logging.Logger):
     def draw(self, xray=True):
         from IPython.display import Image, display
 
-        return display(Image(self.runnable.draw(xray)))
+        image = self.runnable.draw(xray)
+        return display(Image(self.runnable.draw(xray))) if image is not None else None
 
     def draw_to_file(self, filename: str, xray=True):
         from PIL import Image as PILImage

flo_ai/models/flo_agent.py

@@ -7,7 +7,7 @@
 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.state.flo_data_collector import FloDataCollector
+from flo_ai.state.flo_output_collector import FloOutputCollector
 from flo_ai.parsers.flo_parser import FloParser
 
 
@@ -18,7 +18,7 @@ def __init__(
         agent: Runnable,
         executor: AgentExecutor,
         model_name: str,
-        data_collector: Optional[FloDataCollector] = None,
+        data_collector: Optional[FloOutputCollector] = None,
     ) -> None:
         super().__init__(name, executor, ExecutableType.agentic)
         self.model_name = model_name
@@ -36,7 +36,7 @@ def create(
         on_error: Union[str, Callable] = True,
         llm: Union[BaseLanguageModel, None] = None,
         parser: Optional[FloParser] = None,
-        data_collector: Optional[FloDataCollector] = None,
+        data_collector: Optional[FloOutputCollector] = None,
     ):
         model_name = 'default' if llm is None else llm.name
         return FloAgent.Builder(
@@ -65,7 +65,7 @@ def __init__(
             on_error: Union[str, Callable] = True,
             model_name: Union[str, None] = 'default',
             parser: Optional[FloParser] = None,
-            data_collector: Optional[FloDataCollector] = None,
+            data_collector: Optional[FloOutputCollector] = None,
         ) -> None:
             prompt: Union[ChatPromptTemplate, str] = job
             self.name: str = name

flo_ai/models/flo_llm_agent.py

@@ -7,7 +7,7 @@
 from langchain_core.output_parsers import StrOutputParser
 from flo_ai.models.flo_executable import ExecutableType
 from flo_ai.parsers.flo_parser import FloParser
-from flo_ai.state.flo_data_collector import FloDataCollector
+from flo_ai.state.flo_output_collector import FloOutputCollector
 
 
 class FloLLMAgent(ExecutableFlo):
@@ -16,7 +16,7 @@ def __init__(
         name: str,
         executor: Runnable,
         model_name: str,
-        data_collector: Optional[FloDataCollector] = None,
+        data_collector: Optional[FloOutputCollector] = None,
     ) -> None:
         super().__init__(name, executor, ExecutableType.llm)
         self.executor: Runnable = executor
@@ -31,7 +31,7 @@ def create(
         role: Optional[str] = None,
         llm: Union[BaseLanguageModel, None] = None,
         parser: Optional[FloParser] = None,
-        data_collector: Optional[FloDataCollector] = None,
+        data_collector: Optional[FloOutputCollector] = None,
     ):
         model_name = 'default' if llm is None else llm.name
         return FloLLMAgent.Builder(
@@ -55,7 +55,7 @@ def __init__(
             llm: Union[BaseLanguageModel, None] = None,
             model_name: str = None,
             parser: Optional[FloParser] = None,
-            data_collector: Optional[FloDataCollector] = None,
+            data_collector: Optional[FloOutputCollector] = None,
         ) -> None:
             self.model_name = model_name
             prompt: Union[ChatPromptTemplate, str] = job