flock-core 0.2.1__py3-none-any.whl → 0.2.3__py3-none-any.whl

flock/core/mixin/dspy_integration.py

@@ -1,26 +1,30 @@
 """Mixin class for integrating with the dspy library."""
 
 import sys
-from typing import Any
-
-import dspy
+from typing import Any, Literal
 
 from flock.core.logging.logging import get_logger
 from flock.core.util.input_resolver import split_top_level
 
 logger = get_logger("flock")
 
+AgentType = (
+    Literal["ReAct"] | Literal["Completion"] | Literal["ChainOfThought"] | None
+)
+
 
 class DSPyIntegrationMixin:
     """Mixin class for integrating with the dspy library."""
 
     def create_dspy_signature_class(
         self, agent_name, description_spec, fields_spec
-    ) -> dspy.Signature:
+    ) -> Any:
         """Trying to create a dynamic class using dspy library."""
         # ---------------------------
         # 1. Parse the class specification.
         # ---------------------------
+        import dspy
+
         base_class = dspy.Signature
 
         # Start building the class dictionary with a docstring and annotations dict.
@@ -116,13 +120,16 @@ class DSPyIntegrationMixin:
         return type("dspy_" + agent_name, (base_class,), class_dict)
 
     def _configure_language_model(self) -> None:
+        import dspy
+
         """Initialize and configure the language model using dspy."""
         lm = dspy.LM(self.model, cache=self.use_cache)
         dspy.configure(lm=lm)
 
     def _select_task(
         self,
-        signature: dspy.Signature,
+        signature: Any,
+        agent_type_override: AgentType,
     ) -> Any:
         """Select and instantiate the appropriate task based on tool availability.
 
@@ -134,17 +141,36 @@ class DSPyIntegrationMixin:
         Returns:
             An instance of a dspy task (either ReAct or Predict).
         """
+        import dspy
+
         dspy_solver = None
-        if self.tools:
-            dspy_solver = dspy.ReAct(
-                signature,
-                tools=self.tools,
-                max_iters=10,
-            )
+
+        if agent_type_override:
+            if agent_type_override == "ChainOfThought":
+                dspy_solver = dspy.ChainOfThought(
+                    signature,
+                )
+            if agent_type_override == "ReAct":
+                dspy.ReAct(
+                    signature,
+                    tools=self.tools,
+                    max_iters=10,
+                )
+            if agent_type_override == "Completion":
+                dspy_solver = dspy.Predict(
+                    signature,
+                )
         else:
-            dspy_solver = dspy.Predict(
-                signature,
-            )
+            if self.tools:
+                dspy_solver = dspy.ReAct(
+                    signature,
+                    tools=self.tools,
+                    max_iters=10,
+                )
+            else:
+                dspy_solver = dspy.Predict(
+                    signature,
+                )
 
         return dspy_solver
 

flock/core/registry/agent_registry.py

@@ -1,8 +1,14 @@
-"""Registry for storing and managing agents and tools."""
+"""Registry for storing and managing agents and tools with logging and tracing integration."""
 
 from collections.abc import Callable
 
+from opentelemetry import trace
+
 from flock.core.flock_agent import FlockAgent
+from flock.core.logging.logging import get_logger
+
+logger = get_logger("registry")
+tracer = trace.get_tracer(__name__)
 
 
 class Registry:
@@ -16,88 +22,97 @@ class Registry:
     _instance = None
 
     def __new__(cls):
-        """Singleton pattern implementation."""
-        if cls._instance is None:
-            cls._instance = super().__new__(cls)
-            cls._instance._initialize()
-        return cls._instance
+        with tracer.start_as_current_span("Registry.__new__") as span:
+            if cls._instance is None:
+                cls._instance = super().__new__(cls)
+                cls._instance._initialize()
+                logger.info("Registry instance created")
+                span.set_attribute("instance.created", True)
+            return cls._instance
 
     def _initialize(self):
-        """Initialize the registry's storage."""
-        self._agents: list[FlockAgent] = []
-        self._tools: list[tuple[str, Callable]] = []
+        with tracer.start_as_current_span("Registry._initialize"):
+            self._agents: list[FlockAgent] = []
+            self._tools: list[tuple[str, Callable]] = []
+            logger.info("Registry initialized", agents_count=0, tools_count=0)
 
     def register_tool(self, tool_name: str, tool: Callable) -> None:
-        """Register a tool with the registry.
-
-        Args:
-            tool_name: The name to register the tool under
-            tool: The tool function to register
-        """
-        try:
-            self._tools.append((tool_name, tool))
-        except Exception:
-            raise
+        with tracer.start_as_current_span("Registry.register_tool") as span:
+            span.set_attribute("tool_name", tool_name)
+            try:
+                self._tools.append((tool_name, tool))
+                logger.info("Tool registered", tool_name=tool_name)
+            except Exception as e:
+                logger.error(
+                    "Error registering tool", tool_name=tool_name, error=str(e)
+                )
+                span.record_exception(e)
+                raise
 
     def register_agent(self, agent: FlockAgent) -> None:
-        """Register an agent with the registry.
-
-        Args:
-            agent: The agent instance to register
-        """
-        try:
-            self._agents.append(agent)
-        except Exception:
-            raise
+        with tracer.start_as_current_span("Registry.register_agent") as span:
+            span.set_attribute("agent_name", agent.name)
+            try:
+                self._agents.append(agent)
+                logger.info("Agent registered", agent=agent.name)
+            except Exception as e:
+                logger.error(
+                    "Error registering agent", agent=agent.name, error=str(e)
+                )
+                span.record_exception(e)
+                raise
 
     def get_agent(self, name: str) -> FlockAgent | None:
-        """Retrieve an agent by name.
-
-        Args:
-            name: The name of the agent to retrieve
-
-        Returns:
-            The agent instance if found, None otherwise
-        """
-        try:
-            for agent in self._agents:
-                if agent.name == name:
-                    return agent
-            return None
-        except Exception:
-            raise
+        with tracer.start_as_current_span("Registry.get_agent") as span:
+            span.set_attribute("search_agent_name", name)
+            try:
+                for agent in self._agents:
+                    if agent.name == name:
+                        logger.info("Agent found", agent=name)
+                        span.set_attribute("found", True)
+                        return agent
+                logger.warning("Agent not found", agent=name)
+                span.set_attribute("found", False)
+                return None
+            except Exception as e:
+                logger.error("Error retrieving agent", agent=name, error=str(e))
+                span.record_exception(e)
+                raise
 
     def get_tool(self, name: str) -> Callable | None:
-        """Retrieve a tool by name.
-
-        Args:
-            name: The name of the tool to retrieve
-
-        Returns:
-            The tool function if found, None otherwise
-        """
-        try:
-            for tool_name, tool in self._tools:
-                if tool_name == name:
-                    return tool
-            return None
-        except Exception:
-            raise
+        with tracer.start_as_current_span("Registry.get_tool") as span:
+            span.set_attribute("search_tool_name", name)
+            try:
+                for tool_name, tool in self._tools:
+                    if tool_name == name:
+                        logger.info("Tool found", tool=name)
+                        span.set_attribute("found", True)
+                        return tool
+                logger.warning("Tool not found", tool=name)
+                span.set_attribute("found", False)
+                return None
+            except Exception as e:
+                logger.error("Error retrieving tool", tool=name, error=str(e))
+                span.record_exception(e)
+                raise
 
     def get_tools(self, names: list[str] | None) -> list[Callable]:
-        """Retrieve multiple tools by name.
-
-        Args:
-            names: List of tool names to retrieve
-
-        Returns:
-            List of found tool functions (may be empty if none found)
-        """
-        try:
-            if not names:
-                return []
-
-            tools = [self.get_tool(name) for name in names]
-            return [tool for tool in tools if tool is not None]
-        except Exception:
-            raise
+        with tracer.start_as_current_span("Registry.get_tools") as span:
+            span.set_attribute("search_tool_names", str(names))
+            try:
+                if not names:
+                    logger.info("No tool names provided")
+                    return []
+                tools = [self.get_tool(name) for name in names]
+                found_tools = [tool for tool in tools if tool is not None]
+                logger.info(
+                    "Tools retrieved", requested=names, found=len(found_tools)
+                )
+                span.set_attribute("found_tools_count", len(found_tools))
+                return found_tools
+            except Exception as e:
+                logger.error(
+                    "Error retrieving tools", names=str(names), error=str(e)
+                )
+                span.record_exception(e)
+                raise

flock/core/tools/basic_tools.py

@@ -3,7 +3,10 @@
 import importlib
 import os
 
+from flock.core.logging.trace_and_logged import traced_and_logged
 
+
+@traced_and_logged
 def web_search_tavily(query: str):
     """Performs a web search using the Tavily client.
 
@@ -32,9 +35,12 @@ def web_search_tavily(query: str):
         except Exception:
             raise
     else:
-        raise ImportError("Optional tool dependencies not installed. Install with 'pip install flock-core[tools]'.")
+        raise ImportError(
+            "Optional tool dependencies not installed. Install with 'pip install flock-core[tools]'."
+        )
 
 
+@traced_and_logged
 def get_web_content_as_markdown(url: str):
     """Fetches web content from a URL and converts it to Markdown.
 
@@ -51,7 +57,10 @@ def get_web_content_as_markdown(url: str):
         ImportError: If either 'httpx' or 'markdownify' dependencies are not installed.
         Exception: Re-raises any exceptions encountered during the HTTP request or conversion.
     """
-    if importlib.util.find_spec("httpx") is not None and importlib.util.find_spec("markdownify") is not None:
+    if (
+        importlib.util.find_spec("httpx") is not None
+        and importlib.util.find_spec("markdownify") is not None
+    ):
         import httpx
         from markdownify import markdownify as md
 
@@ -63,9 +72,12 @@ def get_web_content_as_markdown(url: str):
         except Exception:
             raise
     else:
-        raise ImportError("Optional tool dependencies not installed. Install with 'pip install flock-core[tools]'.")
+        raise ImportError(
+            "Optional tool dependencies not installed. Install with 'pip install flock-core[tools]'."
+        )
 
 
+@traced_and_logged
 def get_anything_as_markdown(url_or_file_path: str):
     """Converts the content of a URL or file into Markdown format.
 
@@ -94,9 +106,12 @@ def get_anything_as_markdown(url_or_file_path: str):
         except Exception:
             raise
     else:
-        raise ImportError("Optional tool dependencies not installed. Install with 'pip install flock-core[all-tools]'.")
+        raise ImportError(
+            "Optional tool dependencies not installed. Install with 'pip install flock-core[all-tools]'."
+        )
 
 
+@traced_and_logged
 def evaluate_math(expression: str) -> float:
     """Evaluates a mathematical expression using the dspy interpreter.
 
@@ -121,6 +136,7 @@ def evaluate_math(expression: str) -> float:
         raise
 
 
+@traced_and_logged
 def code_eval(python_code: str) -> float:
     """Executes Python code using the dspy interpreter.
 
@@ -145,6 +161,7 @@ def code_eval(python_code: str) -> float:
         raise
 
 
+@traced_and_logged
 def get_current_time() -> str:
     """Retrieves the current time in ISO 8601 format.
 
@@ -157,6 +174,7 @@ def get_current_time() -> str:
     return time
 
 
+@traced_and_logged
 def count_words(text: str) -> int:
     """Counts the number of words in the provided text.
 
@@ -172,6 +190,7 @@ def count_words(text: str) -> int:
     return count
 
 
+@traced_and_logged
 def extract_urls(text: str) -> list[str]:
     """Extracts all URLs from the given text.
 
@@ -191,6 +210,7 @@ def extract_urls(text: str) -> list[str]:
     return urls
 
 
+@traced_and_logged
 def extract_numbers(text: str) -> list[float]:
     """Extracts all numerical values from the provided text.
 
@@ -209,6 +229,7 @@ def extract_numbers(text: str) -> list[float]:
     return numbers
 
 
+@traced_and_logged
 def json_parse_safe(text: str) -> dict:
     """Safely parses a JSON string into a dictionary.
 
@@ -230,6 +251,7 @@ def json_parse_safe(text: str) -> dict:
         return {}
 
 
+@traced_and_logged
 def save_to_file(content: str, filename: str):
     """Saves the given content to a file.
 
@@ -250,6 +272,7 @@ def save_to_file(content: str, filename: str):
         raise
 
 
+@traced_and_logged
 def read_from_file(filename: str) -> str:
     """Reads and returns the content of a file.
 

flock/core/tools/dev_tools/github.py

@@ -5,7 +5,10 @@ import os
 
 import httpx
 
+from flock.core.logging.trace_and_logged import traced_and_logged
 
+
+@traced_and_logged
 def create_user_stories_as_github_issue(title: str, body: str) -> str:
     """Create a new GitHub issue representing a user story.
 
@@ -44,6 +47,7 @@ def create_user_stories_as_github_issue(title: str, body: str) -> str:
         return "Failed to create issue. Please try again later."
 
 
+@traced_and_logged
 def upload_readme(content: str):
     """Upload or update the README.md file in a GitHub repository.
 
@@ -66,7 +70,9 @@ def upload_readme(content: str):
     GITHUB_TOKEN = os.getenv("GITHUB_PAT")
 
     if not GITHUB_USERNAME or not REPO_NAME or not GITHUB_TOKEN:
-        raise ValueError("Missing environment variables: GITHUB_USERNAME, GITHUB_REPO, or GITHUB_PAT")
+        raise ValueError(
+            "Missing environment variables: GITHUB_USERNAME, GITHUB_REPO, or GITHUB_PAT"
+        )
 
     GITHUB_API_URL = f"https://api.github.com/repos/{GITHUB_USERNAME}/{REPO_NAME}/contents/README.md"
 
@@ -75,13 +81,20 @@ def upload_readme(content: str):
     with httpx.Client() as client:
         response = client.get(
             GITHUB_API_URL,
-            headers={"Authorization": f"token {GITHUB_TOKEN}", "Accept": "application/vnd.github.v3+json"},
+            headers={
+                "Authorization": f"Bearer {GITHUB_TOKEN}",
+                "Accept": "application/vnd.github.v3+json",
+            },
         )
 
         data = response.json()
         sha = data.get("sha", None)
 
-        payload = {"message": "Updating README.md", "content": encoded_content, "branch": "main"}
+        payload = {
+            "message": "Updating README.md",
+            "content": encoded_content,
+            "branch": "main",
+        }
 
         if sha:
             payload["sha"] = sha
@@ -89,7 +102,10 @@ def upload_readme(content: str):
         response = client.put(
             GITHUB_API_URL,
             json=payload,
-            headers={"Authorization": f"token {GITHUB_TOKEN}", "Accept": "application/vnd.github.v3+json"},
+            headers={
+                "Authorization": f"Bearer {GITHUB_TOKEN}",
+                "Accept": "application/vnd.github.v3+json",
+            },
         )
 
         if response.status_code in [200, 201]:
@@ -98,6 +114,7 @@ def upload_readme(content: str):
             print("Failed to upload README.md:", response.json())
 
 
+@traced_and_logged
 def create_files(file_paths) -> str:
     """Create multiple files in a GitHub repository with a predefined content.
 
@@ -122,7 +139,9 @@ def create_files(file_paths) -> str:
         GITHUB_TOKEN = os.getenv("GITHUB_PAT")
 
         if not GITHUB_USERNAME or not REPO_NAME or not GITHUB_TOKEN:
-            raise ValueError("Missing environment variables: GITHUB_USERNAME, GITHUB_REPO, or GITHUB_PAT")
+            raise ValueError(
+                "Missing environment variables: GITHUB_USERNAME, GITHUB_REPO, or GITHUB_PAT"
+            )
 
         encoded_content = base64.b64encode(b"#created by flock").decode()
 
@@ -132,13 +151,20 @@ def create_files(file_paths) -> str:
 
                 response = client.get(
                     GITHUB_API_URL,
-                    headers={"Authorization": f"token {GITHUB_TOKEN}", "Accept": "application/vnd.github.v3+json"},
+                    headers={
+                        "Authorization": f"token {GITHUB_TOKEN}",
+                        "Accept": "application/vnd.github.v3+json",
+                    },
                 )
 
                 data = response.json()
                 sha = data.get("sha", None)
 
-                payload = {"message": f"Creating {file_path}", "content": encoded_content, "branch": "main"}
+                payload = {
+                    "message": f"Creating {file_path}",
+                    "content": encoded_content,
+                    "branch": "main",
+                }
 
                 if sha:
                     print(f"Skipping {file_path}, file already exists.")
@@ -147,7 +173,10 @@ def create_files(file_paths) -> str:
                 response = client.put(
                     GITHUB_API_URL,
                     json=payload,
-                    headers={"Authorization": f"token {GITHUB_TOKEN}", "Accept": "application/vnd.github.v3+json"},
+                    headers={
+                        "Authorization": f"token {GITHUB_TOKEN}",
+                        "Accept": "application/vnd.github.v3+json",
+                    },
                 )
 
                 if response.status_code in [200, 201]:

flock/core/util/cli_helper.py

@@ -8,9 +8,13 @@ def display_banner():
     """Display the Flock banner."""
     banner_text = Text(
         """
-▒█▀▀▀ █░░ █▀▀█ █▀▀ █░█
-▒█▀▀▀ █░░ █░░█ █░░ █▀▄
-▒█░░░ ▀▀▀ ▀▀▀▀ ▀▀▀ ▀░▀
+        🦆   🐓   🐤    🦆
+    ╭━━━━━━━━━━━━━━━━━━━━━━━━╮
+🐓  │ ▒█▀▀▀ █░░ █▀▀█ █▀▀ █░█ │  🐧
+    │ ▒█▀▀▀ █░░ █░░█ █░░ █▀▄ │
+🐧  │ ▒█░░░ ▀▀▀ ▀▀▀▀ ▀▀▀ ▀░▀ │  🐓
+    ╰━━━━━━━━━━━━━━━━━━━━━━━━╯
+        🦆     🐤    🐧     🐓
 """,
         justify="center",
         style="bold orange3",