added local llm class.

Author: JDOxygen

CHANGELOG_zh.md

@@ -13,6 +13,7 @@
 
 ### Added
 - 新增Prompt管理界面，支持在线修改Prompt
+- 支持本地大模型 LocalLLM类
 
 ### Changed
 - 修改es包，由elasticsearch[async] 改为 elasticsearch

config.json

@@ -17,9 +17,7 @@
             "is_detailed_tool_call": true,
             "is_detailed_observation": true
         },
-        "llm": {    
-            "cls": "oxygent.llms.OllamaLLM",
-            "base_url": "http://localhost:11434",
+        "llm": {
             "temperature": 0.1, 
             "max_tokens": 4096, 
             "top_p": 1

oxygent/config.py

@@ -53,8 +53,6 @@ class Config:
             "is_detailed_observation": True,
         },
         "llm": {
-            "cls": "oxygent.llms.OllamaLLM",
-            "base_url": "http://localhost:11434",
             "temperature": 0.1,
             "max_tokens": 4096,
             "top_p": 1,

oxygent/oxy/__init__.py

@@ -16,7 +16,7 @@
 )
 from .function_tools.function_hub import FunctionHub
 from .function_tools.function_tool import FunctionTool
-from .llms import HttpLLM, MockLLM, OpenAILLM
+from .llms import HttpLLM, LocalLLM, MockLLM, OpenAILLM
 from .mcp_tools import MCPTool, SSEMCPClient, StdioMCPClient, StreamableMCPClient
 
 __all__ = [
@@ -31,6 +31,7 @@
     "HttpLLM",
     "OpenAILLM",
     "MockLLM",
+    "LocalLLM",
     "MCPTool",
     "StdioMCPClient",
     "StreamableMCPClient",

oxygent/oxy/llms/__init__.py

@@ -1,9 +1,11 @@
 from .http_llm import HttpLLM
+from .local_llm import LocalLLM
 from .mock_llm import MockLLM
 from .openai_llm import OpenAILLM
 
 __all__ = [
     "HttpLLM",
     "OpenAILLM",
     "MockLLM",
+    "LocalLLM",
 ]

oxygent/oxy/llms/http_llm.py

@@ -77,24 +77,12 @@ async def _execute(self, oxy_request: OxyRequest) -> OxyResponse:
                 if k != "messages":
                     payload[k] = v
         else:
-            llm_config = {
-                k: v
-                for k, v in Config.get_llm_config().items()
-                if k
-                not in {
-                    "cls",
-                    "base_url",
-                    "api_key",
-                    "name",
-                    "model_name",
-                }
-            }
             payload = {
                 "messages": await self._get_messages(oxy_request),
                 "model": self.model_name,
                 "stream": True,
             }
-            payload.update(llm_config)
+            payload.update(Config.get_llm_config())
             for k, v in self.llm_params.items():
                 payload[k] = v
             for k, v in oxy_request.arguments.items():

oxygent/oxy/llms/local_llm.py

@@ -0,0 +1,58 @@
+import logging
+
+from pydantic import Field
+
+from ...config import Config
+from ...schemas import OxyRequest, OxyResponse, OxyState
+from .base_llm import BaseLLM
+
+logger = logging.getLogger(__name__)
+
+
+class LocalLLM(BaseLLM):
+    model_path: str = Field("")
+    device: str = Field("auto")
+    model_name: str = Field("")
+
+    async def init(self):
+        try:
+            import torch
+            from transformers import AutoModelForCausalLM, AutoTokenizer
+        except ImportError as e:
+            raise ImportError(
+                "LocalLLM requires 'torch' and 'transformers' packages."
+                "Please install them using 'pip install torch transformers einops transformers_stream_generator accelerate'"
+            ) from e
+
+        await super().init()
+        # Load model directly
+        self._model = AutoModelForCausalLM.from_pretrained(
+            self.model_path, device_map=self.device, torch_dtype=torch.bfloat16
+        )
+        self._tokenizer = AutoTokenizer.from_pretrained(self.model_path)
+
+    async def _execute(self, oxy_request: OxyRequest) -> OxyResponse:
+        payload = {"model": self.model_name, "stream": False}
+        payload.update(Config.get_llm_config())
+        for k, v in self.llm_params.items():
+            payload[k] = v
+        for k, v in oxy_request.arguments.items():
+            if k == "messages":
+                continue
+            payload[k] = v
+        payload = {"max_new_tokens": 512}
+
+        messages = oxy_request.arguments["messages"]
+
+        input_text = self._tokenizer.apply_chat_template(
+            messages, tokenize=False, add_generation_prompt=True
+        )
+        input_ids = self._tokenizer.encode(input_text, return_tensors="pt")
+        input_ids = input_ids.to(self._model.device)
+        outputs = self._model.generate(input_ids=input_ids, **payload)[0]
+        outputs = outputs[len(input_ids[0]) :]
+
+        return OxyResponse(
+            state=OxyState.COMPLETED,
+            output=self._tokenizer.decode(outputs, skip_special_tokens=True),
+        )

oxygent/oxy/llms/openai_llm.py

@@ -38,24 +38,12 @@ async def _execute(self, oxy_request: OxyRequest) -> OxyResponse:
             OxyResponse: The response containing the model's output with COMPLETED state.
         """
         # Construct payload for OpenAI API request
-        llm_config = {
-            k: v
-            for k, v in Config.get_llm_config().items()
-            if k
-            not in {
-                "cls",
-                "base_url",
-                "api_key",
-                "name",
-                "model_name",
-            }
-        }
         payload = {
             "messages": await self._get_messages(oxy_request),
             "model": self.model_name,
             "stream": True,
         }
-        payload.update(llm_config)
+        payload.update(Config.get_llm_config())
         for k, v in self.llm_params.items():
             payload[k] = v
         for k, v in oxy_request.arguments.items():