|
|
import os |
|
|
from dataclasses import dataclass, field |
|
|
from datetime import timedelta |
|
|
from typing import Any, Dict, List, Optional, Tuple |
|
|
|
|
|
from injector import Module, provider |
|
|
|
|
|
from taskweaver.config.module_config import ModuleConfig |
|
|
from taskweaver.misc.component_registry import ComponentRegistry |
|
|
from taskweaver.utils import read_yaml, validate_yaml |
|
|
|
|
|
|
|
|
@dataclass |
|
|
class PluginParameter: |
|
|
"""PluginParameter is the data structure for plugin parameters (including arguments and return values.)""" |
|
|
|
|
|
name: str = "" |
|
|
type: str = "None" |
|
|
required: bool = False |
|
|
description: Optional[str] = None |
|
|
|
|
|
@staticmethod |
|
|
def from_dict(d: Dict[str, Any]): |
|
|
return PluginParameter( |
|
|
name=d["name"], |
|
|
description=d["description"], |
|
|
required=d["required"] if "required" in d else False, |
|
|
type=d["type"] if "type" in d else "Any", |
|
|
) |
|
|
|
|
|
def format_prompt(self, indent: int = 0) -> str: |
|
|
lines: List[str] = [] |
|
|
|
|
|
def line(cnt: str): |
|
|
lines.append(" " * indent + cnt) |
|
|
|
|
|
line(f"- name: {self.name}") |
|
|
line(f" type: {self.type}") |
|
|
line(f" required: {self.required}") |
|
|
line(f" description: {self.description}") |
|
|
|
|
|
return "\n".join(lines) |
|
|
|
|
|
|
|
|
@dataclass |
|
|
class PluginSpec: |
|
|
"""PluginSpec is the data structure for plugin specification defined in the yaml files.""" |
|
|
|
|
|
name: str = "" |
|
|
description: str = "" |
|
|
args: List[PluginParameter] = field(default_factory=list) |
|
|
returns: List[PluginParameter] = field(default_factory=list) |
|
|
embedding: List[float] = field(default_factory=list) |
|
|
|
|
|
@staticmethod |
|
|
def from_dict(d: Dict[str, Any]): |
|
|
return PluginSpec( |
|
|
name=d["name"], |
|
|
description=d["description"], |
|
|
args=[PluginParameter.from_dict(p) for p in d["parameters"]], |
|
|
returns=[PluginParameter.from_dict(p) for p in d["returns"]], |
|
|
embedding=[], |
|
|
) |
|
|
|
|
|
def format_prompt(self) -> str: |
|
|
def normalize_type(t: str) -> str: |
|
|
if t.lower() == "string": |
|
|
return "str" |
|
|
if t.lower() == "integer": |
|
|
return "int" |
|
|
return t |
|
|
|
|
|
def normalize_description(d: str) -> str: |
|
|
d = d.strip().replace("\n", "\n# ") |
|
|
return d |
|
|
|
|
|
def normalize_value(v: PluginParameter) -> PluginParameter: |
|
|
return PluginParameter( |
|
|
name=v.name, |
|
|
type=normalize_type(v.type), |
|
|
required=v.required, |
|
|
description=normalize_description(v.description or ""), |
|
|
) |
|
|
|
|
|
def format_arg_val(val: PluginParameter) -> str: |
|
|
val = normalize_value(val) |
|
|
type_val = f"Optional[{val.type}]" if val.type != "Any" and not val.required else "Any" |
|
|
if val.description is not None: |
|
|
return f"\n# {val.description}\n{val.name}: {type_val}" |
|
|
return f"{val.name}: {type_val}" |
|
|
|
|
|
param_list = ",".join([format_arg_val(p) for p in self.args]) |
|
|
|
|
|
return_type = "" |
|
|
if len(self.returns) > 1: |
|
|
|
|
|
def format_return_val(val: PluginParameter) -> str: |
|
|
val = normalize_value(val) |
|
|
if val.description is not None: |
|
|
return f"\n# {val.name}: {val.description}\n{val.type}" |
|
|
return val.type |
|
|
|
|
|
return_type = f"Tuple[{','.join([format_return_val(r) for r in self.returns])}]" |
|
|
elif len(self.returns) == 1: |
|
|
rv = normalize_value(self.returns[0]) |
|
|
if rv.description is not None: |
|
|
return_type = f"\\\n# {rv.name}: {rv.description}\n{rv.type}" |
|
|
return_type = rv.type |
|
|
else: |
|
|
return_type = "None" |
|
|
return f"# {self.description}\ndef {self.name}({param_list}) -> {return_type}:...\n" |
|
|
|
|
|
|
|
|
@dataclass |
|
|
class PluginEntry: |
|
|
name: str |
|
|
plugin_only: bool |
|
|
impl: str |
|
|
spec: PluginSpec |
|
|
config: Dict[str, Any] |
|
|
required: bool |
|
|
enabled: bool = True |
|
|
|
|
|
@staticmethod |
|
|
def from_yaml_file(path: str) -> Optional["PluginEntry"]: |
|
|
content = read_yaml(path) |
|
|
return PluginEntry.from_yaml_content(content) |
|
|
|
|
|
@staticmethod |
|
|
def from_yaml_content(content: Dict) -> Optional["PluginEntry"]: |
|
|
do_validate = False |
|
|
valid_state = False |
|
|
if do_validate: |
|
|
valid_state = validate_yaml(content, schema="plugin_schema") |
|
|
if not do_validate or valid_state: |
|
|
spec: PluginSpec = PluginSpec.from_dict(content) |
|
|
return PluginEntry( |
|
|
name=spec.name, |
|
|
impl=content.get("code", spec.name), |
|
|
spec=spec, |
|
|
config=content.get("configurations", {}), |
|
|
required=content.get("required", False), |
|
|
enabled=content.get("enabled", True), |
|
|
plugin_only=content.get("plugin_only", False), |
|
|
) |
|
|
return None |
|
|
|
|
|
def format_prompt(self) -> str: |
|
|
return self.spec.format_prompt() |
|
|
|
|
|
def to_dict(self): |
|
|
return { |
|
|
"name": self.name, |
|
|
"impl": self.impl, |
|
|
"spec": self.spec, |
|
|
"config": self.config, |
|
|
"required": self.required, |
|
|
"enabled": self.enabled, |
|
|
} |
|
|
|
|
|
def format_function_calling(self) -> Dict: |
|
|
assert self.plugin_only is True, "Only `plugin_only` plugins can be called in this way." |
|
|
|
|
|
def map_type(t: str) -> str: |
|
|
if t.lower() == "string" or t.lower() == "str" or t.lower() == "text": |
|
|
return "string" |
|
|
if t.lower() == "integer" or t.lower() == "int": |
|
|
return "integer" |
|
|
if t.lower() == "float" or t.lower() == "double" or t.lower() == "number": |
|
|
return "number" |
|
|
if t.lower() == "boolean" or t.lower() == "bool": |
|
|
return "boolean" |
|
|
if t.lower() == "null" or t.lower() == "none": |
|
|
return "null" |
|
|
raise Exception(f"unknown type {t}") |
|
|
|
|
|
function = {"type": "function", "function": {}} |
|
|
required_params = [] |
|
|
function["function"]["name"] = self.name |
|
|
function["function"]["description"] = self.spec.description |
|
|
function["function"]["parameters"] = {"type": "object", "properties": {}} |
|
|
for arg in self.spec.args: |
|
|
function["function"]["parameters"]["properties"][arg.name] = { |
|
|
"type": map_type(arg.type), |
|
|
"description": arg.description, |
|
|
} |
|
|
if arg.required: |
|
|
required_params.append(arg.name) |
|
|
function["function"]["parameters"]["required"] = required_params |
|
|
|
|
|
return function |
|
|
|
|
|
|
|
|
class PluginRegistry(ComponentRegistry[PluginEntry]): |
|
|
def __init__( |
|
|
self, |
|
|
file_glob: str, |
|
|
ttl: Optional[timedelta] = None, |
|
|
) -> None: |
|
|
super().__init__(file_glob, ttl) |
|
|
|
|
|
def _load_component(self, path: str) -> Tuple[str, PluginEntry]: |
|
|
entry: Optional[PluginEntry] = PluginEntry.from_yaml_file(path) |
|
|
if entry is None: |
|
|
raise Exception(f"failed to loading plugin from {path}") |
|
|
if not entry.enabled: |
|
|
raise Exception(f"plugin {entry.name} is disabled") |
|
|
return entry.name, entry |
|
|
|
|
|
|
|
|
class PluginModuleConfig(ModuleConfig): |
|
|
def _configure(self) -> None: |
|
|
self._set_name("plugin") |
|
|
app_dir = self.src.app_base_path |
|
|
self.base_path = self._get_path( |
|
|
"base_path", |
|
|
os.path.join( |
|
|
app_dir, |
|
|
"plugins", |
|
|
), |
|
|
) |
|
|
|
|
|
|
|
|
class PluginModule(Module): |
|
|
@provider |
|
|
def provide_plugin_registry( |
|
|
self, |
|
|
config: PluginModuleConfig, |
|
|
) -> PluginRegistry: |
|
|
import os |
|
|
|
|
|
file_glob = os.path.join(config.base_path, "*.yaml") |
|
|
return PluginRegistry( |
|
|
file_glob=file_glob, |
|
|
ttl=timedelta(minutes=10), |
|
|
) |
|
|
|
