dspex

"""
DSPy gRPC Adapter for Snakepit

This adapter provides gRPC-based DSPy operations for DSPex integration.
It replaces the old JSON-based protocol with efficient gRPC communication.
"""

import os
import json
import logging
import traceback
from typing import Dict, Any, Optional, List
from dataclasses import dataclass

from ..base_adapter import BaseAdapter, tool
from ..session_context import SessionContext
from ..serialization import TypeSerializer

logger = logging.getLogger(__name__)

# Try to import DSPy
try:
    import dspy
    DSPY_AVAILABLE = True
except ImportError:
    DSPY_AVAILABLE = False
    logger.warning("DSPy not available. Some functionality will be limited.")

# Module-level storage for persistence across adapter instances
_MODULE_STORAGE = {}


class DSPyGRPCHandler(BaseAdapter):
    """gRPC adapter for DSPy operations in DSPex"""
    
    def __init__(self):
        super().__init__()
        self.session_context = None
        
    def set_session_context(self, context: SessionContext):
        """Set the session context for variable access"""
        self.session_context = context
        
    @tool(description="Check DSPy availability and version")
    def check_dspy(self) -> Dict[str, Any]:
        """Check if DSPy is available and return version info"""
        if not DSPY_AVAILABLE:
            return {"available": False, "error": "DSPy not installed"}
            
        try:
            version = getattr(dspy, '__version__', 'unknown')
            return {
                "available": True,
                "version": version,
                "modules": dir(dspy)
            }
        except Exception as e:
            return {"available": False, "error": str(e)}
    
    @tool(description="Configure DSPy with a language model")
    def configure_lm(self, model_type: str, **kwargs) -> Dict[str, Any]:
        """Configure DSPy with a language model"""
        if not DSPY_AVAILABLE:
            return {"success": False, "error": "DSPy not available"}
            
        try:
            # Handle different model types
            if model_type == "openai":
                from dspy import OpenAI
                api_key = kwargs.get('api_key', os.getenv('OPENAI_API_KEY'))
                model = kwargs.get('model', 'gpt-3.5-turbo')
                lm = OpenAI(model=model, api_key=api_key)
            elif model_type == "anthropic":
                from dspy import Claude
                api_key = kwargs.get('api_key', os.getenv('ANTHROPIC_API_KEY'))
                model = kwargs.get('model', 'claude-3-sonnet-20240229')
                lm = Claude(model=model, api_key=api_key)
            elif model_type == "gemini":
                # Handle Gemini configuration
                import google.generativeai as genai
                api_key = kwargs.get('api_key', os.getenv('GOOGLE_API_KEY'))
                genai.configure(api_key=api_key)
                # Note: DSPy might need a wrapper for Gemini
                return {"success": True, "message": "Gemini configured (may need DSPy wrapper)"}
            else:
                return {"success": False, "error": f"Unknown model type: {model_type}"}
                
            # Configure DSPy
            dspy.configure(lm=lm)
            
            # Store the LM for later use
            _MODULE_STORAGE['default_lm'] = lm
            
            return {"success": True, "model_type": model_type}
            
        except Exception as e:
            logger.error(f"Error configuring LM: {e}")
            return {"success": False, "error": str(e)}
    
    @tool(description="Create a DSPy signature from specification")
    def create_signature(self, name: str, inputs: List[Dict], outputs: List[Dict], 
                        docstring: Optional[str] = None) -> Dict[str, Any]:
        """Create a DSPy signature"""
        if not DSPY_AVAILABLE:
            return {"success": False, "error": "DSPy not available"}
            
        try:
            # Build signature string
            input_str = ", ".join([f"{field['name']}: {field.get('type', 'str')}" 
                                  for field in inputs])
            output_str = ", ".join([f"{field['name']}: {field.get('type', 'str')}" 
                                   for field in outputs])
            signature_str = f"{input_str} -> {output_str}"
            
            # Create signature class
            class_attrs = {'__doc__': docstring or f"Signature for {name}"}
            
            # Add input fields
            for field in inputs:
                field_type = self._get_dspy_field_type(field.get('type', 'str'))
                field_desc = field.get('description', '')
                class_attrs[field['name']] = dspy.InputField(desc=field_desc)
                
            # Add output fields  
            for field in outputs:
                field_type = self._get_dspy_field_type(field.get('type', 'str'))
                field_desc = field.get('description', '')
                class_attrs[field['name']] = dspy.OutputField(desc=field_desc)
                
            # Create signature class
            signature_class = type(name, (dspy.Signature,), class_attrs)
            
            # Store it
            _MODULE_STORAGE[f"signature_{name}"] = signature_class
            
            return {
                "success": True,
                "name": name,
                "signature_string": signature_str
            }
            
        except Exception as e:
            logger.error(f"Error creating signature: {e}")
            return {"success": False, "error": str(e)}
    
    @tool(description="Create a DSPy program")
    def create_program(self, name: str, program_type: str, signature_name: str,
                      **kwargs) -> Dict[str, Any]:
        """Create a DSPy program"""
        if not DSPY_AVAILABLE:
            return {"success": False, "error": "DSPy not available"}
            
        try:
            # Get signature
            signature = _MODULE_STORAGE.get(f"signature_{signature_name}")
            if not signature:
                # Try simple string signature
                signature = signature_name
                
            # Create program based on type
            if program_type == "predict":
                program = dspy.Predict(signature)
            elif program_type == "chain_of_thought":
                program = dspy.ChainOfThought(signature)
            elif program_type == "react":
                program = dspy.ReAct(signature)
            elif program_type == "program_of_thought":
                program = dspy.ProgramOfThought(signature)
            else:
                return {"success": False, "error": f"Unknown program type: {program_type}"}
                
            # Store program
            _MODULE_STORAGE[name] = program
            
            return {
                "success": True,
                "name": name,
                "type": program_type
            }
            
        except Exception as e:
            logger.error(f"Error creating program: {e}")
            return {"success": False, "error": str(e)}
    
    @tool(description="Execute a DSPy program")
    def execute_program(self, name: str, inputs: Dict[str, Any]) -> Dict[str, Any]:
        """Execute a DSPy program with given inputs"""
        if not DSPY_AVAILABLE:
            return {"success": False, "error": "DSPy not available"}
            
        try:
            program = _MODULE_STORAGE.get(name)
            if not program:
                return {"success": False, "error": f"Program not found: {name}"}
                
            # Execute program
            result = program(**inputs)
            
            # Convert result to dict
            if hasattr(result, 'toDict'):
                output = result.toDict()
            else:
                # Extract non-private attributes
                output = {k: v for k, v in result.__dict__.items() 
                         if not k.startswith('_')}
                         
            return {
                "success": True,
                "result": output
            }
            
        except Exception as e:
            logger.error(f"Error executing program: {e}")
            return {"success": False, "error": str(e), "traceback": traceback.format_exc()}
    
    @tool(description="Store an object for later use")
    def store_object(self, name: str, object_type: str, data: Any) -> Dict[str, Any]:
        """Store an object in the session"""
        try:
            _MODULE_STORAGE[name] = data
            return {"success": True, "name": name}
        except Exception as e:
            return {"success": False, "error": str(e)}
    
    @tool(description="Retrieve a stored object")
    def retrieve_object(self, name: str) -> Dict[str, Any]:
        """Retrieve a stored object"""
        try:
            obj = _MODULE_STORAGE.get(name)
            if obj:
                # Try to serialize the object info
                return {
                    "success": True,
                    "name": name,
                    "type": type(obj).__name__,
                    "info": str(obj)
                }
            else:
                return {"success": False, "error": f"Object not found: {name}"}
        except Exception as e:
            return {"success": False, "error": str(e)}
    
    @tool(description="List stored objects")
    def list_stored_objects(self) -> Dict[str, Any]:
        """List all stored objects"""
        try:
            objects = []
            for name, obj in _MODULE_STORAGE.items():
                objects.append({
                    "name": name,
                    "type": type(obj).__name__
                })
            return {"success": True, "objects": objects}
        except Exception as e:
            return {"success": False, "error": str(e)}
    
    @tool(description="Get current DSPy settings")
    def get_settings(self) -> Dict[str, Any]:
        """Get current DSPy settings"""
        if not DSPY_AVAILABLE:
            return {"error": "DSPy not available"}
        
        try:
            settings = dspy.settings
            settings_dict = {}
            
            # Just get the core settings attributes we care about
            core_attrs = ['lm', 'rm', 'trace', 'explain', 'demonstrate', 'backoff_time', 
                        'branch_idx', 'trace_settings', 'force_rerun', 'compiling']
            
            for attr in core_attrs:
                try:
                    value = getattr(settings, attr, None)
                    
                    # Convert complex objects to string representation
                    if value is None:
                        settings_dict[attr] = None
                    elif isinstance(value, (str, int, float, bool)):
                        settings_dict[attr] = value
                    elif isinstance(value, (list, tuple)):
                        # Convert list items - check each item carefully
                        safe_list = []
                        for item in value:
                            if item is None:
                                safe_list.append(None)
                            elif isinstance(item, (str, int, float, bool)):
                                safe_list.append(item)
                            elif hasattr(item, '__module__'):
                                # Complex object - convert to string
                                safe_list.append(str(item))
                            else:
                                safe_list.append(str(item))
                        settings_dict[attr] = safe_list
                    elif isinstance(value, dict):
                        # Convert dict values - check each value carefully
                        safe_dict = {}
                        for k, v in value.items():
                            if v is None:
                                safe_dict[str(k)] = None
                            elif isinstance(v, (str, int, float, bool)):
                                safe_dict[str(k)] = v
                            elif hasattr(v, '__module__'):
                                # Complex object - convert to string
                                safe_dict[str(k)] = str(v)
                            else:
                                safe_dict[str(k)] = str(v)
                        settings_dict[attr] = safe_dict
                    else:
                        # Complex object - convert to string
                        settings_dict[attr] = str(value)
                except Exception as e:
                    logger.warning(f"Could not serialize settings.{attr}: {e}")
                    settings_dict[attr] = f"<{type(value).__name__} object>" if 'value' in locals() else None
                    
            return settings_dict
            
        except Exception as e:
            logger.error(f"Error getting settings: {e}")
            # Return minimal settings on error
            return {"error": str(e), "lm": None, "rm": None}
    
    @tool(description="Get DSPy statistics")
    def get_stats(self) -> Dict[str, Any]:
        """Get statistics about DSPy usage"""
        try:
            # Count DSPy objects in storage
            programs = [k for k in _MODULE_STORAGE.keys() if not k.startswith('signature_')]
            stats = {
                "dspy_available": DSPY_AVAILABLE,
                "programs_count": len(programs),
                "stored_objects_count": len(_MODULE_STORAGE),
                "programs": programs,
                "has_lm": 'default_lm' in _MODULE_STORAGE
            }
            
            if DSPY_AVAILABLE:
                # Add DSPy specific stats
                try:
                    stats["dspy_version"] = getattr(dspy, '__version__', 'unknown')
                    stats["has_configured_lm"] = hasattr(dspy.settings, 'lm') and dspy.settings.lm is not None
                except:
                    pass
                    
            return {"success": True, "stats": stats}
        except Exception as e:
            return {"success": False, "error": str(e)}
    
    def _get_dspy_field_type(self, type_str: str):
        """Convert type string to DSPy field type"""
        # For now, DSPy mainly uses string types
        # This can be extended as DSPy adds more type support
        return str
    
    # Legacy execute_tool method for compatibility
    def execute_tool(self, tool_name: str, context: SessionContext, arguments: Optional[Dict[str, Any]] = None, **kwargs) -> Any:
        """Legacy tool execution method"""
        self.session_context = context
        # Handle both args and arguments parameters
        args = arguments or kwargs.get('args', {})
        
        # Handle Python.call style invocations
        if "." in tool_name or tool_name.startswith("dspy."):
            return self._handle_python_call(tool_name, args)
        
        # Map old commands to new tool methods
        command_map = {
            "call": lambda: self._handle_call_command(args),
            "check_dspy": self.check_dspy,
            "configure_lm": lambda: self.configure_lm(**args),
            "create_signature": lambda: self.create_signature(**args),
            "create_program": lambda: self.create_program(**args),
            "execute_program": lambda: self.execute_program(**args),
            "store": lambda: self.store_object(**args),
            "retrieve": lambda: self.retrieve_object(**args),
            "list_stored": self.list_stored_objects,
            "get_stats": self.get_stats,
            "get_settings": self.get_settings,  # Add this mapping
            "ping": lambda: self._ping(args),
            "list_programs": self._list_programs,
            "delete_program": lambda: self._delete_program(**args),
            "cleanup": lambda: self._cleanup(**args),
            "reset_state": self._reset_state
        }
        
        if tool_name in command_map:
            return command_map[tool_name]()
        else:
            raise ValueError(f"Unknown tool: {tool_name}")
    
    def _handle_call_command(self, args: Dict[str, Any]) -> Any:
        """Handle the 'call' command from Snakepit.Python.call"""
        target = args.get("target", "")
        kwargs = args.get("kwargs", {})
        store_as = args.get("store_as")
        
        # If kwargs is a string, try to parse it as JSON
        if isinstance(kwargs, str):
            try:
                kwargs = json.loads(kwargs)
            except json.JSONDecodeError:
                # If it's not valid JSON, leave it as is
                pass
        
        result = self._handle_python_call(target, kwargs)
        
        # If store_as is specified, store the actual result object
        if store_as:
            # Check if we have a result object to store
            if "_store_object" in result:
                _MODULE_STORAGE[store_as] = result["_store_object"]
                # Remove _store_object from result to avoid serialization issues
                del result["_store_object"]
                # Store metadata about the object
                result["stored_as"] = store_as
            elif "_metadata" in result and "result" in result:
                # Store the actual object from result
                _MODULE_STORAGE[store_as] = result["result"]
                # Return just the metadata
                return {"result": result["_metadata"], "stored_as": store_as}
            elif "result" in result:
                # Check if result is a DSPy object that needs special handling
                obj = result["result"]
                if hasattr(obj, '__module__') and obj.__module__ and 'dspy' in obj.__module__:
                    # Store the DSPy object and return metadata
                    _MODULE_STORAGE[store_as] = obj
                    return {"result": {"type": type(obj).__name__, "stored_as": store_as}}
                else:
                    # Simple result, just store it
                    _MODULE_STORAGE[store_as] = obj
            
        # If there's metadata, return that instead of the actual object for serialization
        if "_metadata" in result and "result" in result:
            # We have both metadata and an object - return just metadata
            return {"result": result["_metadata"]}
            
        return result
    
    def _handle_python_call(self, target: str, kwargs: Dict[str, Any]) -> Any:
        """Handle Python.call style invocations (e.g. 'dspy.__version__')"""
        logger.info(f"Handling Python.call for target: {target}")
        if not DSPY_AVAILABLE and target.startswith("dspy."):
            raise RuntimeError("DSPy not available")
            
        try:
            # Handle special cases
            if target == "dspy.__version__":
                return {"result": {"value": getattr(dspy, '__version__', 'unknown')}}
            elif target == "dspy.__name__":
                return {"result": {"value": "dspy"}}
            elif target == "dspy.settings":
                # Handle dspy.settings specially - return a dict representation
                try:
                    settings = dspy.settings
                    settings_dict = {}
                    
                    # Just get the core settings attributes we care about
                    core_attrs = ['lm', 'rm', 'trace', 'explain', 'demonstrate', 'backoff_time', 
                                'branch_idx', 'trace_settings', 'force_rerun', 'compiling']
                    
                    for attr in core_attrs:
                        try:
                            value = getattr(settings, attr, None)
                            logger.debug(f"Processing settings.{attr}: type={type(value).__name__ if value else 'None'}")
                            
                            # Convert complex objects to string representation
                            if value is None:
                                settings_dict[attr] = None
                            elif isinstance(value, (str, int, float, bool)):
                                settings_dict[attr] = value
                            elif isinstance(value, (list, tuple)):
                                # Convert list items to strings if they're DSPy objects
                                safe_list = []
                                for i, item in enumerate(value):
                                    logger.debug(f"  List item {i}: type={type(item).__name__}")
                                    # Check if it's a DSPy object
                                    if hasattr(item, '__module__'):
                                        module_name = str(getattr(item, '__module__', ''))
                                        if 'dspy' in module_name or module_name == '__main__':
                                            # Convert DSPy object to string
                                            safe_list.append(str(item))
                                        else:
                                            safe_list.append(item)
                                    else:
                                        safe_list.append(item)
                                settings_dict[attr] = safe_list
                            elif isinstance(value, dict):
                                # Convert dict values to strings if they're DSPy objects
                                safe_dict = {}
                                for k, v in value.items():
                                    logger.debug(f"  Dict key {k}: type={type(v).__name__}")
                                    if hasattr(v, '__module__'):
                                        module_name = str(getattr(v, '__module__', ''))
                                        if 'dspy' in module_name or module_name == '__main__':
                                            safe_dict[str(k)] = str(v)
                                        else:
                                            safe_dict[str(k)] = v
                                    else:
                                        safe_dict[str(k)] = v
                                settings_dict[attr] = safe_dict
                            else:
                                # Complex object - convert to string
                                logger.debug(f"  Converting complex object to string: {type(value).__name__}")
                                settings_dict[attr] = str(value)
                        except Exception as e:
                            logger.warning(f"Could not serialize settings.{attr}: {e}")
                            settings_dict[attr] = f"<{type(value).__name__} object>" if 'value' in locals() else None
                            
                    logger.info(f"Final settings_dict: {list(settings_dict.keys())}")
                    logger.info(f"Settings_dict types: {[(k, type(v).__name__) for k, v in settings_dict.items()]}")
                            
                    # Double-check that we don't have any complex objects
                    clean_dict = {}
                    for k, v in settings_dict.items():
                        if isinstance(v, (str, int, float, bool, type(None))):
                            clean_dict[k] = v
                        elif isinstance(v, (list, tuple)):
                            # Make sure all list items are serializable
                            clean_list = []
                            for item in v:
                                # Check if it's a basic type first
                                if item is None or isinstance(item, (str, int, float, bool)):
                                    clean_list.append(item)
                                else:
                                    # Complex object - convert to string
                                    clean_list.append(str(item))
                            clean_dict[k] = clean_list
                        elif isinstance(v, dict):
                            # Make sure all dict values are serializable
                            clean_subdict = {}
                            for dk, dv in v.items():
                                if hasattr(dv, '__module__'):
                                    clean_subdict[str(dk)] = str(dv)
                                else:
                                    clean_subdict[str(dk)] = dv
                            clean_dict[k] = clean_subdict
                        else:
                            clean_dict[k] = str(v)
                            
                    return {"result": clean_dict}
                    
                except Exception as e:
                    logger.error(f"Error handling dspy.settings: {e}")
                    # Return minimal settings on error
                    return {"result": {"error": str(e), "lm": None, "rm": None}}
            elif target.startswith("stored."):
                # Handle stored object access
                parts = target[7:].split(".", 1)  # Remove "stored." prefix
                obj_name = parts[0]
                obj = _MODULE_STORAGE.get(obj_name)
                if obj:
                    if len(parts) == 1:
                        return {"result": str(obj)}
                    elif parts[1] == "__call__":
                        # Execute the stored program
                        result = obj(**kwargs)
                        if hasattr(result, 'toDict'):
                            return {"result": result.toDict()}
                        else:
                            return {"result": {k: v for k, v in result.__dict__.items() 
                                             if not k.startswith('_')}}
                    else:
                        # Access attribute or call method
                        attr = getattr(obj, parts[1])
                        # Check if it's a callable (method)
                        if callable(attr):
                            # It's a method - call it with kwargs
                            logger.info(f"Calling method {parts[1]} on stored object {obj_name} with kwargs: {kwargs}")
                            # Resolve any stored references in kwargs
                            resolved_kwargs = {}
                            for k, v in kwargs.items():
                                if isinstance(v, str) and v.startswith("stored."):
                                    # Resolve stored reference
                                    stored_name = v[7:]  # Remove "stored." prefix
                                    stored_obj = _MODULE_STORAGE.get(stored_name)
                                    if stored_obj:
                                        resolved_kwargs[k] = stored_obj
                                    else:
                                        raise ValueError(f"Stored object not found: {stored_name}")
                                else:
                                    resolved_kwargs[k] = v
                            
                            result = attr(**resolved_kwargs)
                            # Handle the result based on what it is
                            if hasattr(result, '__module__') and result.__module__ and 'dspy' in result.__module__:
                                # It's a DSPy object - store it and return metadata
                                result_id = f"{obj_name}_{parts[1]}_result"
                                _MODULE_STORAGE[result_id] = result
                                return {"result": {"type": type(result).__name__, "stored_as": result_id}}
                            elif hasattr(result, 'toDict'):
                                return {"result": result.toDict()}
                            else:
                                return {"result": result}
                        else:
                            # It's a regular attribute
                            return {"result": attr}
                else:
                    raise ValueError(f"Stored object not found: {obj_name}")
            else:
                # General Python evaluation (be careful with this in production!)
                # Special handling for dspy.LM
                if target == "dspy.LM":
                    # dspy.LM constructor
                    logger.info(f"Creating dspy.LM with kwargs: {kwargs}")
                    logger.info(f"kwargs type: {type(kwargs)}, content: {kwargs}")
                    # Ensure kwargs is a dict
                    if not isinstance(kwargs, dict):
                        # If it's a string, try to parse it as JSON
                        if isinstance(kwargs, str):
                            try:
                                kwargs = json.loads(kwargs)
                                logger.info(f"Parsed kwargs from JSON: {kwargs}")
                            except json.JSONDecodeError:
                                raise TypeError(f"kwargs must be a dict or valid JSON string, got {type(kwargs)}: {kwargs}")
                        else:
                            raise TypeError(f"kwargs must be a dict, got {type(kwargs)}")
                    lm = dspy.LM(**kwargs)
                    # Return the actual LM object for storage, with metadata for serialization
                    return {"result": lm, "_metadata": {"type": "LM", "model": kwargs.get("model", "unknown")}}
                elif target == "dspy.configure":
                    # Handle dspy.configure - need to resolve stored references
                    if "lm" in kwargs and isinstance(kwargs["lm"], str) and kwargs["lm"].startswith("stored."):
                        # Resolve stored LM reference
                        lm_name = kwargs["lm"][7:]  # Remove "stored." prefix
                        lm = _MODULE_STORAGE.get(lm_name)
                        if lm:
                            kwargs["lm"] = lm
                    dspy.configure(**kwargs)
                    return {"result": {"status": "configured"}}
                elif target in ["dspy.Predict", "dspy.ChainOfThought", "dspy.ReAct", "dspy.ProgramOfThought"]:
                    # Handle DSPy module creation
                    parts = target.split(".")
                    module_class = getattr(dspy, parts[1])
                    module_instance = module_class(**kwargs)
                    # Return metadata for serialization, actual object will be stored
                    return {"result": module_instance, "_metadata": {"type": parts[1], "signature": kwargs.get("signature", "")}}
                elif target == "dspy.Example":
                    # Handle Example creation
                    logger.info(f"Creating dspy.Example with kwargs: {kwargs}")
                    example = dspy.Example(**kwargs)
                    logger.info(f"Created example: {type(example).__name__}")
                    # DSPy Example objects store data as attributes
                    # Extract all non-private attributes
                    example_dict = {}
                    for attr in dir(example):
                        if not attr.startswith('_') and not callable(getattr(example, attr, None)):
                            try:
                                value = getattr(example, attr)
                                example_dict[attr] = value
                            except:
                                pass
                    
                    # If no attributes found, use the kwargs
                    if not example_dict:
                        example_dict = kwargs
                        
                    logger.info(f"Example dict: {example_dict}")
                    # Return dict for serialization, but mark that we have an object to store
                    return {"result": example_dict, "_store_object": example}
                elif target == "dspy.datasets.Dataset":
                    # Handle Dataset creation
                    dataset = dspy.datasets.Dataset(**kwargs)
                    return {"result": {"type": "Dataset", "length": len(dataset)}, "_store_object": dataset}
                else:
                    # General case - handle simple attribute access
                    logger.info(f"General case for target: {target}")
                    parts = target.split(".")
                    obj = globals().get(parts[0])
                    if not obj and parts[0] == "dspy":
                        obj = dspy
                    for part in parts[1:]:
                        obj = getattr(obj, part)
                    if callable(obj):
                        logger.info(f"Calling {target} with kwargs: {kwargs}")
                        result = obj(**kwargs)
                        logger.info(f"Result type: {type(result).__name__}")
                        # Check if result is a DSPy object that can't be serialized
                        if hasattr(result, '__module__') and result.__module__.startswith('dspy'):
                            logger.info(f"Result is a DSPy object, returning with metadata")
                            return {"result": result, "_metadata": {"type": type(result).__name__}}
                        return {"result": result}
                    else:
                        logger.info(f"Returning attribute: {type(obj).__name__}")
                        # Check if it's a complex DSPy object that needs special handling
                        if hasattr(obj, '__module__') and obj.__module__.startswith('dspy'):
                            return {"result": str(obj)}
                        return {"result": obj}
        except Exception as e:
            logger.error(f"Error in Python.call for {target}: {e}")
            raise
    
    def _ping(self, args: Dict[str, Any]) -> Dict[str, Any]:
        """Health check"""
        programs = [k for k in _MODULE_STORAGE.keys() if not k.startswith('signature_')]
        return {
            "status": "ok",
            "dspy_available": DSPY_AVAILABLE,
            "programs_count": len(programs),
            "stored_objects_count": len(_MODULE_STORAGE)
        }
    
    def _list_programs(self) -> Dict[str, Any]:
        """List all programs"""
        programs = [k for k in _MODULE_STORAGE.keys() if not k.startswith('signature_') and not k == 'default_lm']
        return {"programs": programs}
    
    def _delete_program(self, name: str) -> Dict[str, Any]:
        """Delete a program"""
        if name in _MODULE_STORAGE:
            del _MODULE_STORAGE[name]
            return {"success": True}
        return {"success": False, "error": f"Program not found: {name}"}
    
    def _cleanup(self, **kwargs) -> Dict[str, Any]:
        """Cleanup resources"""
        # In gRPC mode, cleanup is handled by session management
        return {"status": "ok"}
    
    def _reset_state(self) -> Dict[str, Any]:
        """Reset all state"""
        _MODULE_STORAGE.clear()
        return {"status": "reset"}