flock-core 0.4.542__py3-none-any.whl → 0.5.0__py3-none-any.whl

flock/core/mcp/mcp_client_manager.py

@@ -1,201 +0,0 @@
-"""Manages a pool of connections for a particular server."""
-
-import copy
-from abc import ABC, abstractmethod
-from asyncio import Lock
-from typing import Any, Generic, TypeVar
-
-from opentelemetry import trace
-from pydantic import (
-    BaseModel,
-    ConfigDict,
-    Field,
-)
-
-from flock.core.logging.logging import get_logger
-from flock.core.mcp.flock_mcp_tool_base import FlockMCPToolBase
-from flock.core.mcp.mcp_client import (
-    FlockMCPClientBase,
-)
-from flock.core.mcp.mcp_config import FlockMCPConfigurationBase
-
-logger = get_logger("mcp.client_manager")
-tracer = trace.get_tracer(__name__)
-
-TClient = TypeVar("TClient", bound="FlockMCPClientBase")
-
-
-class FlockMCPClientManagerBase(BaseModel, ABC, Generic[TClient]):
-    """Handles a Pool of MCPClients of type TClient."""
-
-    client_config: FlockMCPConfigurationBase = Field(
-        ..., description="Configuration for clients."
-    )
-
-    lock: Lock = Field(
-        default_factory=Lock,
-        description="Lock for mutex access.",
-        exclude=True,
-    )
-
-    clients: dict[str, dict[str, FlockMCPClientBase]] = Field(
-        default_factory=dict,
-        exclude=True,
-        description="Internal Store for the clients.",
-    )
-
-    # --- Pydantic v2 Configuratioin ---
-    model_config = ConfigDict(
-        arbitrary_types_allowed=True,
-    )
-
-    @abstractmethod
-    async def make_client(
-        self,
-        additional_params: dict[str, Any] | None = None,
-    ) -> type[TClient]:
-        """Instantiate-but don't connect yet-a fresh client of the concrete subtype."""
-        # default implementation
-        pass
-
-    async def get_client(
-        self,
-        agent_id: str,
-        run_id: str,
-        additional_params: dict[str, Any] | None = None,
-    ) -> type[TClient]:
-        """Provides a client from the pool."""
-        # Attempt to get a client from the client store.
-        # clients are stored like this: agent_id -> run_id -> client
-        with tracer.start_as_current_span("client_manager.get_client") as span:
-            span.set_attribute("agent_id", agent_id)
-            span.set_attribute("run_id", run_id)
-            async with self.lock:
-                try:
-                    logger.debug(
-                        f"Attempting to get client for server '{self.client_config.name}'"
-                    )
-                    refresh = False
-                    if additional_params:
-                        refresh = bool(
-                            additional_params.get("refresh_client", False)
-                        )
-                    client = None
-                    run_clients = self.clients.get(agent_id, None)
-                    if run_clients is None or refresh:
-                        # This means, that across all runs, no agent has ever needed a client.
-                        # This also means that we need to create a client.
-                        client = await self.make_client(
-                            additional_params=copy.deepcopy(additional_params)
-                        )
-                        # Insert the freshly created client
-                        self.clients[agent_id] = {}
-                        self.clients[agent_id][run_id] = client
-
-                    else:
-                        # This means there is at least one entry for the agent_id available
-                        # Now, all we need to do is check if the run_id matches the entrie's run_id
-                        client = run_clients.get(run_id, None)
-                        if client is None or refresh:
-                            # Means no client here with the respective run_id
-                            client = await self.make_client(
-                                additional_params=copy.deepcopy(
-                                    additional_params
-                                )
-                            )
-                            # Insert the freshly created client.
-                            self.clients[agent_id][run_id] = client
-
-                    return client
-                except Exception as e:
-                    # Log the exception and raise it so it becomes visible downstream
-                    logger.error(
-                        f"Unexpected Exception ocurred while trying to get client for server '{self.client_config.name}' with agent_id: {agent_id} and run_id: {run_id}: {e}"
-                    )
-                    span.record_exception(e)
-                    raise e
-
-    async def call_tool(
-        self,
-        agent_id: str,
-        run_id: str,
-        name: str,
-        arguments: dict[str, Any],
-        additional_params: dict[str, Any] | None = None,
-    ) -> Any:
-        """Call a tool."""
-        with tracer.start_as_current_span("client_manager.call_tool") as span:
-            span.set_attribute("agent_id", agent_id)
-            span.set_attribute("run_id", run_id)
-            span.set_attribute("tool_name", name)
-            span.set_attribute("arguments", str(arguments))
-            try:
-                client = await self.get_client(
-                    agent_id=agent_id,
-                    run_id=run_id,
-                    additional_params=additional_params,
-                )
-                result = await client.call_tool(
-                    agent_id=agent_id,
-                    run_id=run_id,
-                    name=name,
-                    arguments=arguments,
-                )
-                return result
-            except Exception as e:
-                logger.error(
-                    f"Exception occurred while trying to call tool {name} on server '{self.client_config.name}': {e}"
-                )
-                span.record_exception(e)
-                return None
-
-    async def get_tools(
-        self,
-        agent_id: str,
-        run_id: str,
-        additional_params: dict[str, Any] | None = None,
-    ) -> list[FlockMCPToolBase]:
-        """Retrieves a list of tools for the agents to act on."""
-        with tracer.start_as_current_span("client_manager.get_tools") as span:
-            span.set_attribute("agent_id", agent_id)
-            span.set_attribute("run_id", run_id)
-            try:
-                client = await self.get_client(
-                    agent_id=agent_id,
-                    run_id=run_id,
-                    additional_params=additional_params,
-                )
-                tools: list[FlockMCPToolBase] = await client.get_tools(
-                    agent_id=agent_id, run_id=run_id
-                )
-                return tools
-            except Exception as e:
-                logger.error(
-                    f"Exception occurred while trying to retrieve Tools for server '{self.client_config.name}' with agent_id: {agent_id} and run_id: {run_id}: {e}"
-                )
-                span.record_exception(e)
-                return []
-
-    async def close_all(self) -> None:
-        """Closes all connections in the pool and cancels background tasks."""
-        with tracer.start_as_current_span("client_manager.close_all") as span:
-            async with self.lock:
-                for agent_id, run_dict in self.clients.items():
-                    logger.debug(
-                        f"Shutting down all clients for agent_id: {agent_id}"
-                    )
-                    for run_id, client in run_dict.items():
-                        logger.debug(
-                            f"Shutting down client for agent_id {agent_id} and run_id {run_id}"
-                        )
-                        try:
-                            await client.disconnect()
-                        except Exception as e:
-                            logger.error(
-                                f"Error when trying to disconnect client for server '{self.client_config.name}': {e}"
-                            )
-                            span.record_exception(e)
-                self.clients = {}  # Let the GC take care of the rest.
-                logger.info(
-                    f"All clients disconnected for server '{self.client_config.name}'"
-                )

flock/core/mcp/types/__init__.py

@@ -1 +0,0 @@
-"""MCP Types package."""

flock/core/mixin/dspy_integration.py

@@ -1,445 +0,0 @@
-# src/flock/core/mixin/dspy_integration.py
-"""Mixin class for integrating with the dspy library."""
-
-import ast
-import re  # Import re for parsing
-import typing
-from typing import Any, Literal
-
-from dspy import Tool
-
-from flock.core.logging.logging import get_logger
-from flock.core.util.splitter import split_top_level
-
-# Import split_top_level (assuming it's moved or copied appropriately)
-# Option 1: If moved to a shared util
-# from flock.core.util.parsing_utils import split_top_level
-# Option 2: If kept within this file (as in previous example)
-# Define split_top_level here or ensure it's imported
-
-logger = get_logger("mixin.dspy")
-
-# Type definition for agent type override
-AgentType = Literal["ReAct", "Completion", "ChainOfThought"] | None
-
-
-# Helper function to resolve type strings (can be static or module-level)
-def _resolve_type_string(type_str: str) -> type:
-    """Resolves a type string into a Python type object.
-    Handles built-ins, registered types, and common typing generics like
-    List, Dict, Optional, Union, Literal.
-    """
-    # Import registry here to avoid circular imports
-    from flock.core.flock_registry import get_registry
-
-    FlockRegistry = get_registry()
-
-    type_str = type_str.strip()
-    logger.debug(f"Attempting to resolve type string: '{type_str}'")
-
-    # 1. Check built-ins and registered types directly
-    try:
-        # This covers str, int, bool, Any, and types registered by name
-        resolved_type = FlockRegistry.get_type(type_str)
-        logger.debug(f"Resolved '{type_str}' via registry to: {resolved_type}")
-        return resolved_type
-    except KeyError:
-        logger.debug(
-            f"'{type_str}' not found directly in registry, attempting generic parsing."
-        )
-        pass  # Not found, continue parsing generics
-
-    # 2. Handle typing generics (List, Dict, Optional, Union, Literal)
-    # Use regex to match pattern like Generic[InnerType1, InnerType2, ...]
-    generic_match = re.fullmatch(r"(\w+)\s*\[(.*)\]", type_str)
-    if generic_match:
-        base_name = generic_match.group(1).strip()
-        args_str = generic_match.group(2).strip()
-        logger.debug(
-            f"Detected generic pattern: Base='{base_name}', Args='{args_str}'"
-        )
-
-        try:
-            # Get the base generic type (e.g., list, dict, Optional) from registry/builtins
-            BaseType = FlockRegistry.get_type(
-                base_name
-            )  # Expects List, Dict etc. to be registered
-            logger.debug(
-                f"Resolved base generic type '{base_name}' to: {BaseType}"
-            )
-
-            # Special handling for Literal
-            if BaseType is typing.Literal:
-                # Split literal values, remove quotes, strip whitespace
-                def parse_literal_args(args_str: str) -> tuple[str, ...]:
-                    try:
-                        return tuple(ast.literal_eval(f"[{args_str}]"))
-                    except (SyntaxError, ValueError) as exc:
-                        raise ValueError(
-                            f"Cannot parse {args_str!r} as literals"
-                        ) from exc
-
-                literal_args = parse_literal_args(args_str)
-                logger.debug(
-                    f"Parsing Literal arguments: {args_str} -> {literal_args}"
-                )
-                resolved_type = typing.Literal[literal_args]  # type: ignore
-                logger.debug(f"Constructed Literal type: {resolved_type}")
-                return resolved_type
-
-            # Recursively resolve arguments for other generics
-            logger.debug(f"Splitting generic arguments: '{args_str}'")
-            arg_strs = split_top_level(args_str)
-            logger.debug(f"Split arguments: {arg_strs}")
-            if not arg_strs:
-                raise ValueError("Generic type has no arguments.")
-
-            resolved_arg_types = tuple(
-                _resolve_type_string(arg) for arg in arg_strs
-            )
-            logger.debug(f"Resolved generic arguments: {resolved_arg_types}")
-
-            # Construct the generic type hint
-            if BaseType is typing.Optional:
-                if len(resolved_arg_types) != 1:
-                    raise ValueError("Optional requires exactly one argument.")
-                # type: ignore
-                resolved_type = typing.Union[resolved_arg_types[0], type(None)]
-                logger.debug(
-                    f"Constructed Optional type as Union: {resolved_type}"
-                )
-                return resolved_type
-            elif BaseType is typing.Union:
-                if not resolved_arg_types:
-                    raise ValueError("Union requires at least one argument.")
-                # type: ignore
-                resolved_type = typing.Union[resolved_arg_types]
-                logger.debug(f"Constructed Union type: {resolved_type}")
-                return resolved_type
-            elif hasattr(
-                BaseType, "__getitem__"
-            ):  # Check if subscriptable (like list, dict, List, Dict)
-                resolved_type = BaseType[resolved_arg_types]  # type: ignore
-                logger.debug(
-                    f"Constructed subscripted generic type: {resolved_type}"
-                )
-                return resolved_type
-            else:
-                # Base type found but cannot be subscripted
-                logger.warning(
-                    f"Base type '{base_name}' found but is not a standard subscriptable generic. Returning base type."
-                )
-                return BaseType
-
-        except (KeyError, ValueError, IndexError, TypeError) as e:
-            logger.warning(
-                f"Failed to parse generic type '{type_str}': {e}. Falling back."
-            )
-            # Fall through to raise KeyError below if base type itself wasn't found or parsing failed
-
-    # 3. If not resolved by now, raise error
-    logger.error(f"Type string '{type_str}' could not be resolved.")
-    raise KeyError(f"Type '{type_str}' could not be resolved.")
-
-
-class DSPyIntegrationMixin:
-    """Mixin class for integrating with the dspy library."""
-
-    def create_dspy_signature_class(
-        self, agent_name, description_spec, fields_spec
-    ) -> Any:
-        """Creates a dynamic DSPy Signature class from string specifications,
-        resolving types using the FlockRegistry.
-        """
-        try:
-            import dspy
-        except ImportError:
-            logger.error(
-                "DSPy library is not installed. Cannot create DSPy signature. "
-                "Install with: pip install dspy-ai"
-            )
-            raise ImportError("DSPy is required for this functionality.")
-
-        base_class = dspy.Signature
-        class_dict = {"__doc__": description_spec, "__annotations__": {}}
-
-        if "->" in fields_spec:
-            inputs_spec, outputs_spec = fields_spec.split("->", 1)
-        else:
-            inputs_spec, outputs_spec = (
-                fields_spec,
-                "",
-            )  # Assume only inputs if no '->'
-
-        def parse_field(field_str):
-            """Parses 'name: type_str | description' using _resolve_type_string."""
-            field_str = field_str.strip()
-            if not field_str:
-                return None
-
-            parts = field_str.split("|", 1)
-            main_part = parts[0].strip()
-            desc = parts[1].strip() if len(parts) > 1 else None
-
-            if ":" in main_part:
-                name, type_str = [s.strip() for s in main_part.split(":", 1)]
-            else:
-                name = main_part
-                type_str = "str"  # Default type
-
-            try:
-                field_type = _resolve_type_string(type_str)
-            except Exception as e:  # Catch resolution errors
-                logger.error(
-                    f"Failed to resolve type '{type_str}' for field '{name}': {e}. Defaulting to str."
-                )
-                field_type = str
-
-            return name, field_type, desc
-
-        def process_fields(fields_string, field_kind):
-            """Process fields and add to class_dict."""
-            if not fields_string or not fields_string.strip():
-                return
-
-            split_fields = split_top_level(fields_string)
-            for field in split_fields:
-                if field.strip():
-                    parsed = parse_field(field)
-                    if not parsed:
-                        continue
-                    name, field_type, desc = parsed
-                    class_dict["__annotations__"][name] = (
-                        field_type  # Use resolved type
-                    )
-
-                    FieldClass = (
-                        dspy.InputField
-                        if field_kind == "input"
-                        else dspy.OutputField
-                    )
-                    # DSPy Fields use 'desc' for description
-                    class_dict[name] = (
-                        FieldClass(desc=desc)
-                        if desc is not None
-                        else FieldClass()
-                    )
-
-        try:
-            process_fields(inputs_spec, "input")
-            process_fields(outputs_spec, "output")
-        except Exception as e:
-            logger.error(
-                f"Error processing fields for DSPy signature '{agent_name}': {e}",
-                exc_info=True,
-            )
-            raise ValueError(
-                f"Could not process fields for signature: {e}"
-            ) from e
-
-        # Create and return the dynamic class
-        try:
-            DynamicSignature = type(
-                "dspy_" + agent_name, (base_class,), class_dict
-            )
-            logger.info(
-                f"Successfully created DSPy Signature: {DynamicSignature.__name__} "
-                f"with fields: {DynamicSignature.__annotations__}"
-            )
-            return DynamicSignature
-        except Exception as e:
-            logger.error(
-                f"Failed to create dynamic type 'dspy_{agent_name}': {e}",
-                exc_info=True,
-            )
-            raise TypeError(f"Could not create DSPy signature type: {e}") from e
-
-    def _configure_language_model(
-        self,
-        model: str | None,
-        use_cache: bool,
-        temperature: float,
-        max_tokens: int,
-    ) -> None:
-        """Initialize and configure the language model using dspy."""
-        if model is None:
-            logger.warning(
-                "No model specified for DSPy configuration. Using DSPy default."
-            )
-            # Rely on DSPy's global default or raise error if none configured
-            # import dspy
-            # if dspy.settings.lm is None:
-            #      raise ValueError("No model specified for agent and no global DSPy LM configured.")
-            return
-
-        try:
-            import dspy
-        except ImportError:
-            logger.error(
-                "DSPy library is not installed. Cannot configure language model."
-            )
-            return  # Or raise
-
-        try:
-            # Ensure 'cache' parameter is handled correctly (might not exist on dspy.LM directly)
-            # DSPy handles caching globally or via specific optimizers typically.
-            # We'll configure the LM without explicit cache control here.
-            lm_instance = dspy.LM(
-                model=model,
-                temperature=temperature,
-                max_tokens=max_tokens,
-                cache=use_cache,
-                # Add other relevant parameters if needed, e.g., API keys via dspy.settings
-            )
-            dspy.settings.configure(lm=lm_instance)
-            logger.info(
-                f"DSPy LM configured with model: {model}, temp: {temperature}, max_tokens: {max_tokens}"
-            )
-            # Note: DSPy caching is usually configured globally, e.g., dspy.settings.configure(cache=...)
-            # or handled by optimizers. Setting `cache=use_cache` on dspy.LM might not be standard.
-        except Exception as e:
-            logger.error(
-                f"Failed to configure DSPy language model '{model}': {e}",
-                exc_info=True,
-            )
-            # We need to raise this exception, otherwise Flock will trundle on until it needs dspy.settings.lm and can't find it.
-            raise
-
-    def _select_task(
-        self,
-        signature: Any,
-        override_evaluator_type: AgentType,
-        max_tool_calls: int = 10,
-        tools: list[Any] | None = None,
-        mcp_tools: list[Any] | None = None,
-        kwargs: dict[str, Any] = {},
-    ) -> Any:
-        """Select and instantiate the appropriate DSPy Program/Module."""
-        try:
-            import dspy
-        except ImportError:
-            logger.error(
-                "DSPy library is not installed. Cannot select DSPy task."
-            )
-            raise ImportError("DSPy is required for this functionality.")
-
-        processed_tools = []
-        if tools:
-            for tool in tools:
-                if callable(tool):  # Basic check
-                    processed_tools.append(tool)
-                # Could add more sophisticated tool wrapping/validation here if needed
-                else:
-                    logger.warning(
-                        f"Item '{tool}' in tools list is not callable, skipping."
-                    )
-
-        processed_mcp_tools = []
-        if mcp_tools:
-            for mcp_tool in mcp_tools:
-                if isinstance(mcp_tool, Tool):  # Basic check
-                    processed_mcp_tools.append(mcp_tool)
-                else:
-                    logger.warning(
-                        f"Item '{mcp_tool}' is not a dspy.primitives.Tool, skipping."
-                    )
-
-        dspy_program = None
-        selected_type = override_evaluator_type
-
-        # Determine type if not overridden
-        if not selected_type:
-            selected_type = (
-                "ReAct" if processed_tools or processed_mcp_tools else "Predict"
-            )  # Default logic
-
-        logger.debug(
-            f"Selecting DSPy program type: {selected_type} (Tools provided: {bool(processed_tools)}) (MCP Tools: {bool(processed_mcp_tools)}"
-        )
-
-        # Merge list of native tools and processed tools.
-        # This makes mcp tools appear as native code functions to the llm of the agent.
-        merged_tools = []
-
-        if processed_tools:
-            merged_tools = merged_tools + processed_tools
-
-        if processed_mcp_tools:
-            merged_tools = merged_tools + processed_mcp_tools
-
-        try:
-            if selected_type == "ChainOfThought":
-                dspy_program = dspy.ChainOfThought(signature, **kwargs)
-            elif selected_type == "ReAct":
-                if not kwargs:
-                    kwargs = {"max_iters": max_tool_calls}
-                dspy_program = dspy.ReAct(
-                    signature, tools=merged_tools or [], **kwargs
-                )
-            elif selected_type == "Predict":  # Default or explicitly Completion
-                dspy_program = dspy.Predict(signature)
-            else:  # Fallback or handle unknown type
-                logger.warning(
-                    f"Unknown or unsupported agent_type_override '{selected_type}'. Defaulting to dspy.Predict."
-                )
-                dspy_program = dspy.Predict(signature)
-
-            logger.info(
-                f"Instantiated DSPy program: {type(dspy_program).__name__}"
-            )
-            return dspy_program
-        except Exception as e:
-            logger.error(
-                f"Failed to instantiate DSPy program of type '{selected_type}': {e}",
-                exc_info=True,
-            )
-            raise RuntimeError(f"Could not create DSPy program: {e}") from e
-
-    def _process_result(
-        self, result: Any, inputs: dict[str, Any]
-    ) -> tuple[dict[str, Any], float, list]:
-        """Convert the DSPy result object to a dictionary."""
-        import dspy
-
-        if result is None:
-            logger.warning("DSPy program returned None result.")
-            return {}
-        try:
-            # DSPy Prediction objects often behave like dicts or have .keys() / items()
-            if hasattr(result, "items") and callable(result.items):
-                output_dict = dict(result.items())
-            elif hasattr(result, "__dict__"):  # Fallback for other object types
-                output_dict = {
-                    k: v
-                    for k, v in result.__dict__.items()
-                    if not k.startswith("_")
-                }
-            else:
-                # If it's already a dict (less common for DSPy results directly)
-                if isinstance(result, dict):
-                    output_dict = result
-                else:  # Final fallback
-                    logger.warning(
-                        f"Could not reliably convert DSPy result of type {type(result)} to dict. Returning as is."
-                    )
-                    output_dict = {"raw_result": result}
-
-            logger.debug(f"Processed DSPy result to dict: {output_dict}")
-            # Optionally merge inputs back if desired (can make result dict large)
-            final_result = {**inputs, **output_dict}
-
-            lm = dspy.settings.get("lm")
-            cost = sum([x["cost"] for x in lm.history if x["cost"] is not None])
-            lm_history = lm.history
-
-            return final_result, cost, lm_history
-
-        except Exception as conv_error:
-            logger.error(
-                f"Failed to process DSPy result into dictionary: {conv_error}",
-                exc_info=True,
-            )
-            return {
-                "error": "Failed to process result",
-                "raw_result": str(result),
-            }