Merge pull request #233 from qiandl2000/v0.2.5

Support deepseek thinking content output

Author: XeonHis

omagent-core/src/omagent_core/models/llms/base.py

@@ -206,15 +206,61 @@ def infer(self, input_list: List[Dict[str, Any]], **kwargs) -> List[T]:
         stm_token_usage = self.stm(self.workflow_instance_id).get('token_usage', defaultdict(int))
 
         def process_stream(self, stream_output):
-            for chunk in stream_output:
-                if chunk.usage is not None:
-                    for key, value in chunk.usage.dict().items():
-                        if key in ["prompt_tokens", "completion_tokens", 'total_tokens']:
-                            if value is not None:
-                                stm_token_usage[key] += value
-                    self.stm(self.workflow_instance_id)['token_usage'] = stm_token_usage
-            
-                yield chunk
+
+            # local deepseek
+            if hasattr(next(stream_output).choices[0].delta, "reasoning_content"):
+                reasoning_flag = False
+                answering_flag = False
+                for chunk in stream_output:
+                    # output reasoning
+                    if chunk.choices[0].delta.reasoning_content is not None:
+                        if not reasoning_flag:
+                            chunk.choices[0].delta.content = "Reasoning:" + chunk.choices[0].delta.reasoning_content
+                            reasoning_flag = True
+                        else:
+                            chunk.choices[0].delta.content = chunk.choices[0].delta.reasoning_content
+                    # output Answering
+                    elif not answering_flag:
+                        chunk.choices[0].delta.content = "Answer:" + chunk.choices[0].delta.content
+                        answering_flag = True
+                    if chunk.usage is not None:
+                        for key, value in chunk.usage.dict().items():
+                            if key in ["prompt_tokens", "completion_tokens", 'total_tokens']:
+                                if value is not None:
+                                    stm_token_usage[key] += value
+                        self.stm(self.workflow_instance_id)['token_usage'] = stm_token_usage
+                    yield chunk
+
+            # ollama deepseek
+            elif "deepseek" in self.llm.model_id:
+                reasoning_flag = False
+                answering_flag = False
+                for chunk in stream_output:
+                    if not reasoning_flag:
+                        chunk.choices[0].delta.content = "Reasoning:" + chunk.choices[0].delta.content
+                        reasoning_flag = True
+                    elif not answering_flag:
+                        if chunk.choices[0].delta.content == "</think>":
+                            chunk.choices[0].delta.content = "Answer:"
+                        answering_flag = True
+                    if chunk.usage is not None:
+                        for key, value in chunk.usage.dict().items():
+                            if key in ["prompt_tokens", "completion_tokens", 'total_tokens']:
+                                if value is not None:
+                                    stm_token_usage[key] += value
+                        self.stm(self.workflow_instance_id)['token_usage'] = stm_token_usage
+                    yield chunk
+
+            # models without reasoning
+            else:
+                for chunk in stream_output:
+                    if chunk.usage is not None:
+                        for key, value in chunk.usage.dict().items():
+                            if key in ["prompt_tokens", "completion_tokens", 'total_tokens']:
+                                if value is not None:
+                                    stm_token_usage[key] += value
+                        self.stm(self.workflow_instance_id)['token_usage'] = stm_token_usage
+                    yield chunk
 
         for prompt in prompts:
             output = self.llm.generate(prompt, **kwargs)

omagent-core/src/omagent_core/tool_system/tools/calculator/calculator.py

@@ -17,9 +17,9 @@
 class Calculator(BaseTool):
     args_schema: ArgSchema = ArgSchema(**ARGSCHEMA)
     description: str = (
-        "Calculator tool for executing all mathematical calculation. Please output a mathematical formula. \
+        "Calculator tool for executing simple mathematical calculation. Please output a mathematical formula. \
     Addition, subtraction, multiplication, and division are represented by +, -, *, and /, respectively. \
-    Exponentiation is denoted by **, and square root is implemented using sqrt(). Parentheses can be used normally. \
+    Exponentiation is denoted by **, and square root is implemented using sqrt(). Other functions are not allowed to be used. \
     Only the constants pi and e are available, representing the mathematical constant pi and Euler's number, respectively."
     )