symbolicai 0.20.2__py3-none-any.whl → 1.0.0__py3-none-any.whl

symai/backend/engines/neurosymbolic/engine_llama_cpp.py

@@ -2,15 +2,14 @@ import asyncio
 import json
 import logging
 from copy import deepcopy
+from typing import Any, ClassVar
 
 import aiohttp
-import httpx
 import nest_asyncio
-import requests
 
 from ....core import Argument
 from ....core_ext import retry
-from ....utils import CustomUserWarning
+from ....utils import UserMessage
 from ...base import Engine
 from ...settings import SYMAI_CONFIG, SYMSERVER_CONFIG
 
@@ -25,14 +24,14 @@ class LlamaCppTokenizer:
 
     @staticmethod
     async def _encode(text: str) -> list[int]:
-        async with aiohttp.ClientSession() as session:
-            async with session.post(f"{LlamaCppTokenizer._server_endpoint}/extras/tokenize", json={
-                "input": text,
-            }) as res:
-                if res.status != 200:
-                    CustomUserWarning(f"Request failed with status code: {res.status}", raise_with=ValueError)
-                res = await res.json()
-                return res['tokens']
+        async with aiohttp.ClientSession() as session, session.post(
+            f"{LlamaCppTokenizer._server_endpoint}/tokenize",
+            json={"content": text},
+        ) as res:
+            if res.status != 200:
+                UserMessage(f"Request failed with status code: {res.status}", raise_with=ValueError)
+            response_json = await res.json()
+            return response_json['tokens']
 
     @staticmethod
     def encode(text: str) -> list[int]:
@@ -45,14 +44,14 @@ class LlamaCppTokenizer:
 
     @staticmethod
     async def _decode(tokens: list[int]) -> str:
-        async with aiohttp.ClientSession() as session:
-            async with session.post(f"{LlamaCppTokenizer._server_endpoint}/extras/detokenize", json={
-                "tokens": tokens,
-            }) as res:
-                if res.status != 200:
-                    CustomUserWarning(f"Request failed with status code: {res.status}", raise_with=ValueError)
-                res = await res.json()
-                return res['text']
+        async with aiohttp.ClientSession() as session, session.post(
+            f"{LlamaCppTokenizer._server_endpoint}/detokenize",
+            json={"tokens": tokens},
+        ) as res:
+            if res.status != 200:
+                UserMessage(f"Request failed with status code: {res.status}", raise_with=ValueError)
+            response_json = await res.json()
+            return response_json['content']
 
     @staticmethod
     def decode(tokens: list[int]) -> str:
@@ -65,7 +64,7 @@ class LlamaCppTokenizer:
 
 
 class LlamaCppEngine(Engine):
-    _retry_params = {
+    _retry_params: ClassVar[dict[str, Any]] = {
         'tries': 5,
         'delay': 2,
         'max_delay': 60,
@@ -73,7 +72,7 @@ class LlamaCppEngine(Engine):
         'jitter': (1, 5),
         'graceful': True
     }
-    _timeout_params = {
+    _timeout_params: ClassVar[dict[str, Any]] = {
         'read': None,
         'connect': None,
     }
@@ -91,7 +90,7 @@ class LlamaCppEngine(Engine):
         if self.id() != 'neurosymbolic':
             return
         if not SYMSERVER_CONFIG.get('online'):
-            CustomUserWarning('You are using the llama.cpp engine, but the server endpoint is not started. Please start the server with `symserver [--args]` or run `symserver --help` to see the available options for this engine.', raise_with=ValueError)
+            UserMessage('You are using the llama.cpp engine, but the server endpoint is not started. Please start the server with `symserver [--args]` or run `symserver --help` to see the available options for this engine.', raise_with=ValueError)
         self.server_endpoint = f"http://{SYMSERVER_CONFIG.get('--host')}:{SYMSERVER_CONFIG.get('--port')}"
         self.tokenizer = LlamaCppTokenizer # backwards compatibility with how we handle tokenization, i.e. self.tokenizer().encode(...)
         self.timeout_params = self._validate_timeout_params(timeout_params)
@@ -112,23 +111,23 @@ class LlamaCppEngine(Engine):
         if 'except_remedy' in kwargs:
             self.except_remedy = kwargs['except_remedy']
 
-    def compute_required_tokens(self, messages) -> int:
+    def compute_required_tokens(self, _messages) -> int:
         #@TODO: quite non-trivial how to handle this with the llama.cpp server
-        CustomUserWarning('Not implemented for llama.cpp!', raise_with=NotImplementedError)
+        UserMessage('Not implemented for llama.cpp!', raise_with=NotImplementedError)
 
-    def compute_remaining_tokens(self, prompts: list) -> int:
+    def compute_remaining_tokens(self, _prompts: list) -> int:
         #@TODO: quite non-trivial how to handle this with the llama.cpp server
-        CustomUserWarning('Not implemented for llama.cpp!', raise_with=NotImplementedError)
+        UserMessage('Not implemented for llama.cpp!', raise_with=NotImplementedError)
 
     def _validate_timeout_params(self, timeout_params):
         if not isinstance(timeout_params, dict):
-            CustomUserWarning("timeout_params must be a dictionary", raise_with=ValueError)
+            UserMessage("timeout_params must be a dictionary", raise_with=ValueError)
         assert all(key in timeout_params for key in ['read', 'connect']), "Available keys: ['read', 'connect']"
         return timeout_params
 
     def _validate_retry_params(self, retry_params):
         if not isinstance(retry_params, dict):
-            CustomUserWarning("retry_params must be a dictionary", raise_with=ValueError)
+            UserMessage("retry_params must be a dictionary", raise_with=ValueError)
         assert all(key in retry_params for key in ['tries', 'delay', 'max_delay', 'backoff', 'jitter', 'graceful']), "Available keys: ['tries', 'delay', 'max_delay', 'backoff', 'jitter', 'graceful']"
         return retry_params
 
@@ -138,7 +137,7 @@ class LlamaCppEngine(Engine):
         try:
             current_loop = asyncio.get_event_loop()
             if current_loop.is_closed():
-                CustomUserWarning("Event loop is closed.", raise_with=RuntimeError)
+                UserMessage("Event loop is closed.", raise_with=RuntimeError)
             return current_loop
         except RuntimeError:
             new_loop = asyncio.new_event_loop()
@@ -148,7 +147,7 @@ class LlamaCppEngine(Engine):
     def _prepare_request_payload(self, argument: Argument) -> dict:
         """Prepares the request payload from the argument."""
         kwargs = argument.kwargs
-        return {
+        payload = {
             "messages": argument.prop.prepared_input,
             "temperature": kwargs.get('temperature', 0.6),
             "frequency_penalty": kwargs.get('frequency_penalty', 0),
@@ -162,12 +161,28 @@ class LlamaCppEngine(Engine):
             "repeat_penalty": kwargs.get('repeat_penalty', 1),
             "logits_bias": kwargs.get('logits_bias'),
             "logprobs": kwargs.get('logprobs', False),
-            "functions": kwargs.get('functions'),
-            "function_call": kwargs.get('function_call'),
             "grammar": kwargs.get('grammar'),
             "response_format": kwargs.get('response_format'),
         }
 
+        model = SYMSERVER_CONFIG.get('-m') or SYMSERVER_CONFIG.get('--model')
+        if model:
+            payload["model"] = model
+
+        tools = kwargs.get('tools')
+        if tools:
+            payload["tools"] = tools
+
+        tool_choice = kwargs.get('tool_choice')
+        if tool_choice is not None:
+            payload["tool_choice"] = tool_choice
+
+        extra_body = kwargs.get('extra_body')
+        if isinstance(extra_body, dict):
+            payload.update(extra_body)
+
+        return payload
+
     async def _arequest(self, payload: dict) -> dict:
         """Makes an async HTTP request to the llama.cpp server."""
         @retry(**self.retry_params)
@@ -176,17 +191,29 @@ class LlamaCppEngine(Engine):
                 sock_connect=self.timeout_params['connect'],
                 sock_read=self.timeout_params['read']
             )
-            async with aiohttp.ClientSession(timeout=timeout) as session:
-                async with session.post(
-                    f"{self.server_endpoint}/v1/chat/completions",
-                    json=payload
-                ) as res:
-                    if res.status != 200:
-                        CustomUserWarning(f"Request failed with status code: {res.status}", raise_with=ValueError)
-                    return await res.json()
+            async with aiohttp.ClientSession(timeout=timeout) as session, session.post(
+                f"{self.server_endpoint}/v1/chat/completions",
+                json=payload
+            ) as res:
+                if res.status != 200:
+                    UserMessage(f"Request failed with status code: {res.status}", raise_with=ValueError)
+                return await res.json()
 
         return await _make_request()
 
+    @classmethod
+    def _extract_thinking(cls, response):
+        """Extract reasoning traces from llama.cpp responses."""
+        if not isinstance(response, dict):
+            return None
+        choices = response.get('choices', [])
+        if not isinstance(choices, list) or not choices:
+            return None
+        for choice in choices:
+            if isinstance(choice, dict) and isinstance(choice.get('message'), dict):
+                return choice['message'].get('reasoning_content')
+        return None
+
     def forward(self, argument):
         payload = self._prepare_request_payload(argument)
 
@@ -196,21 +223,63 @@ class LlamaCppEngine(Engine):
         try:
             res = loop.run_until_complete(self._arequest(payload))
         except Exception as e:
-            CustomUserWarning(f'Error during generation. Caused by: {e}', raise_with=ValueError)
+            UserMessage(f'Error during generation. Caused by: {e}', raise_with=ValueError)
 
         metadata = {'raw_output': res}
 
+        if payload.get('tools'):
+            metadata = self._process_tool_calls(res, metadata)
+
+        thinking = self._extract_thinking(res)
+        if thinking:
+            metadata['thinking'] = thinking
+
         output = [r['message']['content'] for r in res['choices']]
         output = output if isinstance(argument.prop.prepared_input, list) else output[0]
 
         return output, metadata
 
+    @staticmethod
+    def _process_tool_calls(res, metadata):
+        choices = res.get('choices') if isinstance(res, dict) else None
+        if not choices:
+            return metadata
+        hit = False
+        for choice in choices:
+            if not isinstance(choice, dict):
+                continue
+            message = choice.get('message') or {}
+            tool_calls = message.get('tool_calls') or []
+            if not tool_calls:
+                continue
+            for tool_call in tool_calls:
+                if not isinstance(tool_call, dict):
+                    continue
+                function = tool_call.get('function') or {}
+                if hit:
+                    UserMessage("Multiple function calls detected in the response but only the first one will be processed.")
+                    return metadata
+                arguments = function.get('arguments')
+                try:
+                    args_dict = json.loads(arguments) if isinstance(arguments, str) else arguments or {}
+                except json.JSONDecodeError:
+                    args_dict = {}
+                metadata['function_call'] = {
+                    'name': function.get('name'),
+                    'arguments': args_dict or {}
+                }
+                hit = True
+                break
+            if hit:
+                break
+        return metadata
+
     def _prepare_raw_input(self, argument):
         if not argument.prop.processed_input:
-            CustomUserWarning('Need to provide a prompt instruction to the engine if raw_input is enabled.', raise_with=ValueError)
+            UserMessage('Need to provide a prompt instruction to the engine if raw_input is enabled.', raise_with=ValueError)
         value = argument.prop.processed_input
-        if type(value) != list:
-            if type(value) != dict:
+        if not isinstance(value, list):
+            if not isinstance(value, dict):
                 value = {'role': 'user', 'content': str(value)}
             value = [value]
         return value
@@ -242,17 +311,17 @@ class LlamaCppEngine(Engine):
 
         payload = argument.prop.payload
         if argument.prop.payload:
-            user += f"<ADDITIONAL_CONTEXT/>\n{str(payload)}\n\n"
+            user += f"<ADDITIONAL_CONTEXT/>\n{payload!s}\n\n"
 
         examples: list[str] = argument.prop.examples
         if examples and len(examples) > 0:
-            user += f"<EXAMPLES/>\n{str(examples)}\n\n"
+            user += f"<EXAMPLES/>\n{examples!s}\n\n"
 
         if argument.prop.prompt is not None and len(argument.prop.prompt) > 0:
-            user += f"<INSTRUCTION/>\n{str(argument.prop.prompt)}\n\n"
+            user += f"<INSTRUCTION/>\n{argument.prop.prompt!s}\n\n"
 
         if argument.prop.template_suffix:
-            user += f" You will only generate content for the placeholder `{str(argument.prop.template_suffix)}` following the instructions and the provided context information.\n\n"
+            user += f" You will only generate content for the placeholder `{argument.prop.template_suffix!s}` following the instructions and the provided context information.\n\n"
 
         user += str(argument.prop.processed_input)