gateway/modules/chat/methodBase.py

158 lines
No EOL
6.3 KiB
Python

from enum import Enum
from typing import Dict, List, Optional, Any, Literal
from datetime import datetime, UTC
from pydantic import BaseModel, Field
import logging
from functools import wraps
import inspect
logger = logging.getLogger(__name__)
def action(func):
"""Decorator to mark a method as an available action"""
@wraps(func)
async def wrapper(self, parameters: Dict[str, Any], *args, **kwargs):
return await func(self, parameters, *args, **kwargs)
wrapper.is_action = True
return wrapper
class MethodBase:
"""Base class for all methods"""
def __init__(self, serviceCenter: Any):
"""Initialize method with service center"""
self.service = serviceCenter
self.name: str
self.description: str
self.logger = logging.getLogger(f"{__name__}.{self.__class__.__name__}")
@property
def actions(self) -> Dict[str, Dict[str, Any]]:
"""Dynamically collect all actions decorated with @action in the class."""
actions = {}
for attr_name in dir(self):
# Skip the actions property itself to avoid recursion
if attr_name == 'actions':
continue
try:
attr = getattr(self, attr_name)
if callable(attr) and getattr(attr, 'is_action', False):
sig = inspect.signature(attr)
params = {}
for param_name, param in sig.parameters.items():
if param_name not in ['self', 'parameters']:
param_type = param.annotation if param.annotation != param.empty else Any
params[param_name] = {
'type': param_type,
'required': param.default == param.empty,
'description': None,
'default': param.default if param.default != param.empty else None
}
actions[attr_name] = {
'description': attr.__doc__ or '',
'parameters': params,
'method': attr
}
except (AttributeError, RecursionError):
# Skip attributes that cause issues
continue
return actions
def getActionSignature(self, actionName: str) -> str:
"""Get formatted action signature for AI prompt generation (detailed version)"""
if actionName not in self.actions:
return ""
action = self.actions[actionName]
paramList = []
# Extract detailed parameter information from docstring
docstring = action.get('description', '')
paramDescriptions, paramTypes = self._extractParameterDetails(docstring)
for paramName in paramDescriptions:
paramType = paramTypes.get(paramName, 'Any')
paramDesc = paramDescriptions.get(paramName, '')
# Mark required parameters with * if possible (not available from docstring, so omit)
if paramDesc:
paramList.append(f"{paramName}:{paramType} # {paramDesc}")
else:
paramList.append(f"{paramName}:{paramType}")
signature = f"{self.name}.{actionName}"
if paramList:
signature += f"({', '.join(paramList)})"
# Add return type and main description
returnType = "ActionResult"
mainDesc = self._extractMainDescription(docstring)
if mainDesc:
signature += f" -> {returnType} # {mainDesc}"
return signature
def _extractParameterDetails(self, docstring: str):
"""Extract parameter names, types, and descriptions from docstring"""
descriptions = {}
types = {}
if not docstring:
return descriptions, types
lines = docstring.split('\n')
inParameters = False
for line in lines:
line = line.strip()
if 'Parameters:' in line:
inParameters = True
continue
elif inParameters and (line.startswith('Returns:') or line.startswith('Raises:') or line.startswith('Args:')):
break
elif inParameters and line:
# Look for parameter descriptions like "paramName (type): description"
if ':' in line and '(' in line:
parts = line.split(':', 1)
if len(parts) == 2:
paramPart = parts[0].strip()
descPart = parts[1].strip()
# Extract parameter name and type
if '(' in paramPart:
paramName = paramPart.split('(')[0].strip()
paramType = paramPart[paramPart.find('(')+1:paramPart.find(')')].strip()
descriptions[paramName] = descPart
types[paramName] = paramType
# Also handle multi-line descriptions
elif line and not line.startswith('Each document') and not line.startswith('contains'):
if descriptions:
lastParam = list(descriptions.keys())[-1]
descriptions[lastParam] += " " + line
return descriptions, types
def _extractMainDescription(self, docstring: str) -> str:
"""Extract main description from docstring"""
if not docstring:
return ""
lines = docstring.split('\n')
mainDesc = ""
for line in lines:
line = line.strip()
if line and not line.startswith('Parameters:') and not line.startswith('Returns:') and not line.startswith('Raises:'):
mainDesc = line
break
return mainDesc
def _formatType(self, type_annotation) -> str:
"""Format type annotation for display"""
if type_annotation == Any:
return "Any"
elif hasattr(type_annotation, '__name__'):
return type_annotation.__name__
elif hasattr(type_annotation, '_name'):
return type_annotation._name
else:
return str(type_annotation)