agent-api-server 2.1.7__py3-none-any.whl

agent_api_server/dynamic_llm/dynamic_llm.py

@@ -0,0 +1,331 @@
+from __future__ import annotations
+import ast
+import json
+import os
+import threading
+import hashlib
+import logging
+from typing import Dict, Any, Optional, List, AsyncGenerator, Iterator, Union, Type
+from langchain_core.tools import BaseTool
+from langchain_core.language_models import LanguageModelInput
+from langchain_core.runnables.config import RunnableConfig
+from langchain_core.runnables import Runnable, RunnableLambda
+from llm_sdk.llm import ModelInstance, ModelInstanceFactory
+from llm_sdk.model_providers.base import ConfigType
+from langchain_core.messages import BaseMessage, BaseMessageChunk
+from model_manage_client import ModelManageClient
+from pydantic import BaseModel
+from functools import partial
+from contextlib import contextmanager
+
+from agent_api_server.configs import global_config
+from agent_api_server.shared.util_func import set_model_config, get_env
+
+logger = logging.getLogger(__name__)
+
+
+class DynamicLLM(Runnable):
+    _instance_lock = threading.Lock()
+    _llm_cache: Dict[str, ModelInstance] = {}
+    _hash_fn = partial(hashlib.md5, usedforsecurity=False)
+
+    def __init__(self, tool_name: str, config_type: ConfigType = ConfigType.CHAT, use_sys_llm: Optional[bool] = False):
+        super().__init__()
+        self._model: Optional[ModelInstance] = None
+        self._current_config_hash: Optional[str] = None
+        self._tools: Optional[List[BaseTool]] = None
+        self.tool_name = tool_name.lower()
+        self.config_type = config_type
+        self.use_sys_llm = use_sys_llm
+        self._lock = threading.RLock()
+
+    def rerank(self, input_dict: Union[List[BaseMessage], List[str], str],
+               documents: List[str], top_n: Optional[int] = None, score_threshold: Optional[float] = None, config: Optional[RunnableConfig] = None) -> Any:
+        if self.config_type != ConfigType.RERANK:
+            raise ValueError("Only rerank model type support rerank function")
+
+        llm_config = self._get_llm_config_from_context(config or {})
+
+        with self._get_model_instance(llm_config, config) as model:
+            return model.rerank(input_dict, documents, top_n, score_threshold)
+
+    def invoke(self, input_dict: Union[List[BaseMessage], List[str], str],
+               config: Optional[RunnableConfig] = None, **kwargs) -> Any:
+        llm_config = self._get_llm_config_from_context(config or {})
+
+        with self._get_model_instance(llm_config, config) as model:
+            if self.config_type == ConfigType.CHAT:
+                if not isinstance(input_dict, list) or not all(isinstance(x, BaseMessage) for x in input_dict):
+                    raise ValueError("Chat model requires List[BaseMessage] input_dict")
+                return model.invoke(input_dict, **kwargs)
+            elif self.config_type == ConfigType.EMBEDDING:
+                if isinstance(input_dict, str):
+                    return model.embed_query(input_dict)
+                return model.embed_documents(input_dict if isinstance(input_dict, list) else [input_dict])
+            elif self.config_type == ConfigType.RERANK:
+                raise ValueError("Rerank model should use rerank() method")
+            else:
+                raise ValueError(f"Unsupported config type: {self.config_type}")
+
+    async def ainvoke(self, input_dict: Union[List[BaseMessage], List[str], str],
+                      config: Optional[RunnableConfig] = None, **kwargs) -> Any:
+        llm_config = self._get_llm_config_from_context(config or {})
+
+        with self._get_model_instance(llm_config, config) as model:
+            if self.config_type == ConfigType.CHAT:
+                if not isinstance(input_dict, list) or not all(isinstance(x, BaseMessage) for x in input_dict):
+                    raise ValueError("Chat model requires List[BaseMessage] input_dict")
+                return await model.ainvoke(input_dict, **kwargs)
+            elif self.config_type == ConfigType.EMBEDDING:
+                if isinstance(input_dict, str):
+                    return await model.aembed_query(input_dict)
+                return await model.aembed_documents(input_dict if isinstance(input_dict, list) else [input_dict])
+            elif self.config_type == ConfigType.RERANK:
+                raise ValueError("Rerank model should use rerank() method")
+            else:
+                raise ValueError(f"Unsupported config type: {self.config_type}")
+
+    def stream(self, input_dict: List[BaseMessage], config: Optional[RunnableConfig] = None,
+               **kwargs) -> Iterator[BaseMessageChunk]:
+        if self.config_type != ConfigType.CHAT:
+            raise ValueError("Only chat models support streaming")
+
+        llm_config = self._get_llm_config_from_context(config or {})
+
+        with self._get_model_instance(llm_config, config) as model:
+            return model.stream(input_dict, **kwargs)
+
+    async def astream(self, input_dict: List[BaseMessage], config: Optional[RunnableConfig] = None,
+                      **kwargs) -> AsyncGenerator[BaseMessageChunk, None]:
+        if self.config_type != ConfigType.CHAT:
+            raise ValueError("Only chat models support streaming")
+
+        llm_config = self._get_llm_config_from_context(config or {})
+
+        with self._get_model_instance(llm_config, config) as model:
+            async for chunk in model.astream(input_dict, **kwargs):
+                yield chunk
+
+    def embed_documents(self, texts: List[str], config: Optional[RunnableConfig] = None) -> List[List[float]]:
+        if self.config_type != ConfigType.EMBEDDING:
+            raise ValueError("Only embedding models support this method")
+
+        llm_config = self._get_llm_config_from_context(config or {})
+
+        with self._get_model_instance(llm_config, config) as model:
+            return model.embed_documents(texts)
+
+    async def aembed_documents(self, texts: List[str], config: Optional[RunnableConfig] = None) -> List[List[float]]:
+        if self.config_type != ConfigType.EMBEDDING:
+            raise ValueError("Only embedding models support this method")
+
+        llm_config = self._get_llm_config_from_context(config or {})
+
+        with self._get_model_instance(llm_config, config) as model:
+            return await model.aembed_documents(texts)
+
+    def embed_query(self, text: str, config: Optional[RunnableConfig] = None) -> List[float]:
+        if self.config_type != ConfigType.EMBEDDING:
+            raise ValueError("Only embedding models support this method")
+
+        llm_config = self._get_llm_config_from_context(config or {})
+
+        with self._get_model_instance(llm_config, config) as model:
+            return model.embed_query(text)
+
+    async def aembed_query(self, text: str, config: Optional[RunnableConfig] = None) -> List[float]:
+        if self.config_type != ConfigType.EMBEDDING:
+            raise ValueError("Only embedding models support this method")
+
+        llm_config = self._get_llm_config_from_context(config or {})
+
+        with self._get_model_instance(llm_config, config) as model:
+            return await model.aembed_query(text)
+
+    def bind_tools(self, tools: List[BaseTool], **kwargs) -> 'DynamicLLM':
+        if self.config_type != ConfigType.CHAT:
+            raise ValueError("Cannot bind tools to non-chat models")
+
+        with self._lock:
+            self._tools = tools
+            if self._model is not None:
+                self._model.bind_tools(tools, **kwargs)
+        return self
+
+    def with_structured_output(self, schema: Union[Dict, Type[BaseModel]], **kwargs) -> Runnable:
+        if self.config_type != ConfigType.CHAT:
+            raise ValueError("Only chat models support structured output")
+
+        def wrapper(input_data: LanguageModelInput, config: Optional[RunnableConfig] = None):
+            llm_config = self._get_llm_config_from_context(config or {})
+            with self._get_model_instance(llm_config, config) as model:
+                return model.with_structured_output(
+                    schema=schema,
+                    **kwargs
+                ).invoke(input_data)
+
+        return RunnableLambda(wrapper)
+
+    @contextmanager
+    def _get_model_instance(self, llm_config: Dict[str, Any], config: RunnableConfig):
+        if self.use_sys_llm:
+            logger.info(f"use_sys_llm is {self.use_sys_llm}, so call model with system llm configuration, CLIENT_TOKEN is {global_config.CLIENT_TOKEN}")
+
+            self._model = ModelInstanceFactory.get_model_instance(
+                model_managment_url=os.environ['MODEL_MANAGER_SERVICE_URL'],
+                config_type=self.config_type,
+                token=os.environ['CLIENT_TOKEN'],
+                app_id=llm_config.get('agent_id'),
+                app_name=llm_config.get('agent_id'),
+            )
+            yield self._model
+        else:
+            if self._is_credentials_empty(llm_config):
+                agent_id = llm_config.get('agent_id')
+                ts_id = llm_config.get('ts_tenant')
+
+                logger.warning(
+                    f"Local cache does not contain LLM credentials for tenant '{ts_id}'. "
+                    f"Retrieve the credentials from Model Management and configure them in the system.")
+
+                m_client =  ModelManageClient(global_config.MODEL_MANAGER_SERVICE_URL, global_config.CLIENT_TOKEN)
+                model_info = m_client.get_agent_models(agent_id, ts_id)
+
+                logger.info(f"LLM credentials for tenant {ts_id} is {model_info}")
+
+                for tool_name, tool_cfg in model_info.items():
+                    for model_type, model_cfg in tool_cfg.items():
+                        model_provider = model_cfg.get("provider", "")
+                        model_name = model_cfg.get("model", "")
+                        model_type = model_type
+                        credentials_json = json.dumps(model_cfg.get("credentials", {}), ensure_ascii=False)
+
+                        logger.info(f"""
+                                       - Tenant ID: {agent_id}
+                                       - Tool: {tool_name}
+                                       - Provider: {model_provider}
+                                       - Model: {model_name}
+                                       - Model Type: {model_type}
+                                       """)
+
+                        set_model_config(
+                            tool_name=tool_name,
+                            model_type=model_type,
+                            model_name=model_name,
+                            model_provider=model_provider,
+                            credentials=credentials_json,
+                            agent_id=agent_id,
+                            ts_id=ts_id
+                        )
+
+                configurable = config.get("configurable", {})
+                configurable.update(dict(get_env(ts_tenant=ts_id)))
+                llm_config = self._get_llm_config_from_context(config or {})
+                with self._lock:
+                    self._ensure_model_initialized(llm_config)
+                    yield self._model
+            else:
+                with self._lock:
+                    self._ensure_model_initialized(llm_config)
+                    yield self._model
+
+    def _ensure_model_initialized(self, config: Dict[str, Any]) -> None:
+        config_hash = self._generate_config_hash(config)
+
+        if self._model is None or self._current_config_hash != config_hash:
+            provider = config.get('provider', 'unknown')
+            cache_key = f"{provider}:{self.config_type.value}:{config_hash}"
+
+            with self._instance_lock:
+                if cache_key in self._llm_cache:
+                    self._model = self._llm_cache[cache_key]
+                else:
+                    credentials = config.get('credentials')
+                    credentials_dict = {}
+
+                    if credentials:
+                        if isinstance(credentials, dict):
+                            credentials_dict = credentials
+                        elif isinstance(credentials, str):
+                            try:
+                                credentials_dict = json.loads(credentials)
+                            except json.JSONDecodeError:
+                                try:
+                                    credentials_dict = ast.literal_eval(credentials)
+                                except (ValueError, SyntaxError) as e:
+                                    logger.error(f"Failed to parse credentials string: {e}")
+                        else:
+                            logger.warning(f"Unexpected credentials type: {type(credentials)}")
+
+                    model_params = {
+                        'provider': config.get('provider'),
+                        'model': config.get('model'),
+                        'credentials': credentials_dict,
+                    }
+                    self._model = ModelInstance(
+                        config_type=self.config_type,
+                        model_params=model_params,
+                        app_id=config.get('agent_id')
+                    )
+                    self._llm_cache[cache_key] = self._model
+
+                if self.config_type == ConfigType.CHAT and self._tools:
+                    self._model.bind_tools(self._tools)
+
+                self._current_config_hash = config_hash
+
+    def _get_llm_config_from_context(self, context: Dict[str, Any]) -> Dict[str, Any]:
+        configurable = context.get("configurable", {})
+        ts_tenant = configurable.get("ts_tenant") or configurable.get("TSTenant")
+        graph_name = configurable.get("graph_name")
+
+        value = configurable.get("use_sys_llm")
+        if value:
+            self.use_sys_llm = True if value == "true" else False
+
+        config = {
+            "provider": configurable.get(self._get_config_key("PROVIDER", ts_tenant)),
+            "model": configurable.get(self._get_config_key("MODEL", ts_tenant)),
+            "credentials": configurable.get(self._get_config_key("CREDENTIALS", ts_tenant)),
+            "agent_id": graph_name,
+            "ts_tenant": ts_tenant
+        }
+
+        return config
+
+    def _get_config_key(self, field: str, ts_tenant: Optional[str] = None) -> str:
+        parts = [self.tool_name.upper(), self.config_type.value, field] if self.tool_name != "default" else [
+            self.config_type.value, field]
+        if ts_tenant:
+            parts.append(ts_tenant)
+        return "_".join(parts)
+
+    @classmethod
+    def _generate_config_hash(cls, config: Dict[str, Any]) -> str:
+        relevant_keys = ['provider', 'model', 'credentials']
+        hash_input = "|".join(str(config.get(k, "")) for k in relevant_keys)
+        return cls._hash_fn(hash_input.encode()).hexdigest()
+
+    @staticmethod
+    def _is_credentials_empty(config: Optional[Dict[str, Any]]) -> bool:
+        if config is None:
+            return True
+
+        credentials = config.get('credentials')
+        if not credentials:
+            return True
+
+        return False
+
+    def clear_cache(self):
+        with self._instance_lock:
+            self._llm_cache.clear()
+            self._model = None
+            self._current_config_hash = None
+
+    @classmethod
+    def get_cache_size(cls) -> int:
+        return len(cls._llm_cache)
+
+    def __repr__(self) -> str:
+        return f"DynamicLLM(tool_name={self.tool_name}, config_type={self.config_type})"