symbolicai 0.21.0__py3-none-any.whl → 1.1.0__py3-none-any.whl

symai/backend/engines/index/engine_vectordb.py

@@ -1,9 +1,10 @@
 import itertools
 from copy import deepcopy
+from typing import ClassVar
 
 from ....extended.vectordb import VectorDB
 from ....symbol import Result
-from ....utils import CustomUserWarning
+from ....utils import UserMessage
 from ...base import Engine
 from ...settings import SYMAI_CONFIG
 
@@ -31,72 +32,76 @@ class VectorDBResult(Result):
         try:
             res = self._to_symbol(res).ast()
         except Exception as e:
-            message = ['Sorry, failed to interact with index. Please check index name and try again later:', str(e)]
+            message = [
+                "Sorry, failed to interact with index. Please check index name and try again later:",
+                str(e),
+            ]
             # package the message for the IndexResult class
-            res = {'matches': [{'metadata': {'text': '\n'.join(message)}}]}
-        return [v['metadata']['text'] for v in res]
+            res = {"matches": [{"metadata": {"text": "\n".join(message)}}]}
+        return [v["metadata"]["text"] for v in res]
 
     def _unpack_matches(self):
         if not self.value:
             return
         for i, match in enumerate(self.value):
-            match = match.strip()
-            if match.startswith('# ----[FILE_START]') and '# ----[FILE_END]' in match:
-                m = match.split('[FILE_CONTENT]:')[-1].strip()
-                splits = m.split('# ----[FILE_END]')
-                assert len(splits) >= 2, 'Invalid file format: {}'.format(splits)
+            match_value = match.strip()
+            if match_value.startswith("# ----[FILE_START]") and "# ----[FILE_END]" in match_value:
+                m = match_value.split("[FILE_CONTENT]:")[-1].strip()
+                splits = m.split("# ----[FILE_END]")
+                assert len(splits) >= 2, f"Invalid file format: {splits}"
                 content = splits[0]
-                file_name = ','.join(splits[1:]) # TODO: check why there are multiple file names
+                file_name = ",".join(splits[1:])  # TODO: check why there are multiple file names
                 yield file_name.strip(), content.strip()
             else:
-                yield i+1, match
+                yield i + 1, match_value
 
     def __str__(self):
-        str_view = ''
+        str_view = ""
         for filename, content in self._unpack_matches():
             # indent each line of the content
-            content = '\n'.join(['  ' + line for line in content.split('\n')])
-            str_view += f'* {filename}\n{content}\n\n'
-        return f'''
+            content_view = "\n".join(["  " + line for line in content.split("\n")])
+            str_view += f"* {filename}\n{content_view}\n\n"
+        return f"""
 [RESULT]
-{'-=-' * 13}
+{"-=-" * 13}
 
 Query: {self._query}
 
-{'-=-' * 13}
+{"-=-" * 13}
 
 Matches:
 
 {str_view}
-{'-=-' * 13}
-'''
+{"-=-" * 13}
+"""
 
     def _repr_html_(self) -> str:
         # return a nicely styled HTML list results based on retrieved documents
-        doc_str = ''
+        doc_str = ""
         for filename, content in self._unpack_matches():
             doc_str += f'<li><a href="{filename}"><b>{filename}</a></b><br>{content}</li>\n'
-        return f'<ul>{doc_str}</ul>'
+        return f"<ul>{doc_str}</ul>"
 
 
 class VectorDBIndexEngine(Engine):
     # Updated default values to be congruent with VectorDB's defaults
-    _default_index_name = 'dataindex'
+    _default_index_name = "dataindex"
     _default_index_dims = 768
     _default_index_top_k = 5
-    _default_index_metric = 'cosine'
-    _index_dict = {}
-    _index_storage_file = None
+    _default_index_metric = "cosine"
+    _index_dict: ClassVar[dict[str, object]] = {}
+    _index_storage_file: ClassVar[str | None] = None
+
     def __init__(
-            self,
-            index_name=_default_index_name,
-            index_dims=_default_index_dims,
-            index_top_k=_default_index_top_k,
-            index_metric=_default_index_metric,
-            index_dict=_index_dict,
-            index_storage_file=_index_storage_file,
-            **kwargs
-        ):
+        self,
+        index_name=_default_index_name,
+        index_dims=_default_index_dims,
+        index_top_k=_default_index_top_k,
+        index_metric=_default_index_metric,
+        index_dict=_index_dict,
+        index_storage_file=_index_storage_file,
+        **_kwargs,
+    ):
         super().__init__()
         self.config = deepcopy(SYMAI_CONFIG)
         self.index_name = index_name
@@ -111,9 +116,12 @@ class VectorDBIndexEngine(Engine):
         self.name = self.__class__.__name__
 
     def id(self) -> str:
-        if not self.config['INDEXING_ENGINE_API_KEY'] or self.config['INDEXING_ENGINE_API_KEY'] == '':
-            return 'index'
-        return super().id() # default to unregistered
+        if (
+            not self.config["INDEXING_ENGINE_API_KEY"]
+            or self.config["INDEXING_ENGINE_API_KEY"] == ""
+        ):
+            return "index"
+        return super().id()  # default to unregistered
 
     def forward(self, argument):
         query = argument.prop.prepared_input
@@ -129,34 +137,51 @@ class VectorDBIndexEngine(Engine):
 
         self._init(index_name, top_k, index_dims, metric)
 
-        if operation == 'search':
+        if operation == "search":
             if isinstance(query, list) and len(query) > 1:
-                CustomUserWarning('VectorDB indexing engine does not support multiple queries. Pass a single string query instead.', raise_with=ValueError)
+                UserMessage(
+                    "VectorDB indexing engine does not support multiple queries. Pass a single string query instead.",
+                    raise_with=ValueError,
+                )
             query_vector = self.index[index_name].embedding_function([query])[0]
-            results = self.index[index_name](vector=query_vector, top_k=top_k, return_similarities=similarities)
-            rsp = [{'metadata': {'text': result}} for result in results]
-        elif operation == 'add':
-            assert isinstance(query, list), 'VectorDB indexing requires a list of queries at insertion, even if there is only one query.'
+            results = self.index[index_name](
+                vector=query_vector, top_k=top_k, return_similarities=similarities
+            )
+            rsp = [{"metadata": {"text": result}} for result in results]
+        elif operation == "add":
+            assert isinstance(query, list), (
+                "VectorDB indexing requires a list of queries at insertion, even if there is only one query."
+            )
             documents = []
             vectors = []
             for q in query:
                 vectors.append(self.index[index_name].embedding_function([q])[0])
                 documents.append(q)
             self.index[index_name].add(documents=documents, vectors=vectors)
-        elif operation == 'config':
-            assert kwargs, 'Please provide a configuration by passing the appropriate kwargs. Currently, only `load`, `save`, `purge`.'
-            maybe_as_prompt = kwargs.get('prompt')
-            if kwargs.get('load', maybe_as_prompt == 'load'):
-                assert storage_file, 'Please provide a `storage_file` path to load the pre-computed index.'
+        elif operation == "config":
+            assert kwargs, (
+                "Please provide a configuration by passing the appropriate kwargs. Currently, only `load`, `save`, `purge`."
+            )
+            maybe_as_prompt = kwargs.get("prompt")
+            if kwargs.get("load", maybe_as_prompt == "load"):
+                assert storage_file, (
+                    "Please provide a `storage_file` path to load the pre-computed index."
+                )
                 self.load(index_name, storage_file, index_dims, top_k, metric)
-            elif kwargs.get('save', maybe_as_prompt == 'save'):
+            elif kwargs.get("save", maybe_as_prompt == "save"):
                 self.save(index_name, storage_file)
-            elif kwargs.get('purge', maybe_as_prompt == 'purge'):
+            elif kwargs.get("purge", maybe_as_prompt == "purge"):
                 self.purge(index_name)
             else:
-                CustomUserWarning('Invalid configuration; please use either "load", "save", or "purge".', raise_with=ValueError)
+                UserMessage(
+                    'Invalid configuration; please use either "load", "save", or "purge".',
+                    raise_with=ValueError,
+                )
         else:
-            CustomUserWarning('Invalid operation; please use either "search", "add", or "config".', raise_with=ValueError)
+            UserMessage(
+                'Invalid operation; please use either "search", "add", or "config".',
+                raise_with=ValueError,
+            )
 
         metadata = {}
         rsp = VectorDBResult(rsp, query[0], None)
@@ -170,12 +195,15 @@ class VectorDBIndexEngine(Engine):
                 index_dims=index_dims,
                 top_k=top_k,
                 similarity_metric=metric,
-                embedding_model=embedding_model #@NOTE: the VectorDBIndexEngine class uses precomputed embeddings so the model is not needed in the VectorDB class
+                embedding_model=embedding_model,  # @NOTE: the VectorDBIndexEngine class uses precomputed embeddings so the model is not needed in the VectorDB class
             )
 
     def prepare(self, argument):
-        assert not argument.prop.processed_input, 'VectorDB indexing engine does not support processed_input.'
+        assert not argument.prop.processed_input, (
+            "VectorDB indexing engine does not support processed_input."
+        )
         argument.prop.prepared_input = argument.prop.prompt
+        argument.prop.limit = 1
 
     def load(self, index_name, storage_file, index_dims, top_k, metric):
         self.index[index_name] = VectorDB(
@@ -186,11 +214,11 @@ class VectorDBIndexEngine(Engine):
             index_name=index_name,
         )
 
-    def save(self, index_name = None, storage_file = None):
+    def save(self, index_name=None, storage_file=None):
         index_name = index_name or self.index_name
         storage_file = storage_file or self._index_storage_file
         self.index[index_name].save(storage_file)
 
-    def purge(self, index_name = None):
+    def purge(self, index_name=None):
         index_name = index_name or self.index_name
         self.index[index_name].purge(index_name)

symai/backend/engines/lean/engine_lean4.py

@@ -1,11 +1,15 @@
-import os
 import subprocess
+import tempfile
+from pathlib import Path
+from typing import Any
+
 import docker
 import paramiko
-import tempfile
-from typing import Any, List, Tuple, Optional, Dict
-from ...base import Engine
+
 from ....symbol import Result
+from ....utils import UserMessage
+from ...base import Engine
+
 
 class LeanResult(Result):
     """
@@ -14,7 +18,8 @@ class LeanResult(Result):
     Attributes:
         _value (Dict[str, str]): A dictionary containing the output of the Lean execution.
     """
-    def __init__(self, value: Dict[str, str]) -> None:
+
+    def __init__(self, value: dict[str, str]) -> None:
         """
         Initializes a new LeanResult instance.
 
@@ -24,6 +29,7 @@ class LeanResult(Result):
         super().__init__(value)
         self._value = value
 
+
 class LeanEngine(Engine):
     """
     Engine for executing Lean code within a Docker container, providing SSH access for execution.
@@ -39,10 +45,10 @@ class LeanEngine(Engine):
 
     def __init__(
         self,
-        ssh_host: str = 'localhost',
+        ssh_host: str = "localhost",
         ssh_port: int = 2222,
-        ssh_user: str = 'root',
-        ssh_key_path: str = '~/.ssh/id_rsa'
+        ssh_user: str = "root",
+        ssh_key_path: str = "~/.ssh/id_rsa",
     ) -> None:
         """
         Initializes the LeanEngine with SSH and Docker configurations.
@@ -57,7 +63,7 @@ class LeanEngine(Engine):
         self.ssh_host: str = ssh_host
         self.ssh_port: int = ssh_port
         self.ssh_user: str = ssh_user
-        self.ssh_key_path: str = os.path.expanduser(ssh_key_path)
+        self.ssh_key_path: Path = Path(ssh_key_path).expanduser()
         self.docker_client: docker.DockerClient = docker.from_env()
         self.container: docker.models.containers.Container = self._ensure_container()
         self.name = self.__class__.__name__
@@ -70,7 +76,7 @@ class LeanEngine(Engine):
         Returns:
             str: The identifier of the LeanEngine, 'lean4'.
         """
-        return 'lean4'
+        return "lean4"
 
     def _ensure_container(self) -> docker.models.containers.Container:
         """
@@ -82,10 +88,14 @@ class LeanEngine(Engine):
         container_name: str = "lean-container"
 
         try:
-            existing_container: docker.models.containers.Container = self.docker_client.containers.get(container_name)
+            existing_container: docker.models.containers.Container = (
+                self.docker_client.containers.get(container_name)
+            )
             existing_container.remove(force=True)
         except docker.errors.NotFound:
-            print(f"No existing container named '{container_name}' found. Proceeding to create a new one.")
+            UserMessage(
+                f"No existing container named '{container_name}' found. Proceeding to create a new one."
+            )
 
         dockerfile: str = """
         FROM buildpack-deps:buster
@@ -110,17 +120,18 @@ class LeanEngine(Engine):
         """
         with tempfile.NamedTemporaryFile("w", delete=False) as temp_dockerfile:
             temp_dockerfile.write(dockerfile)
-            dockerfile_path: str = temp_dockerfile.name
+            dockerfile_path = Path(temp_dockerfile.name)
 
         image: docker.models.images.Image
-        image, _ = self.docker_client.images.build(path=os.path.dirname(dockerfile_path), dockerfile=dockerfile_path, tag="lean4-container-image")
-        os.remove(dockerfile_path)
+        image, _ = self.docker_client.images.build(
+            path=str(dockerfile_path.parent),
+            dockerfile=str(dockerfile_path),
+            tag="lean4-container-image",
+        )
+        dockerfile_path.unlink()
 
         container: docker.models.containers.Container = self.docker_client.containers.run(
-            image.id,
-            detach=True,
-            name=container_name,
-            ports={'22/tcp': self.ssh_port}
+            image.id, detach=True, name=container_name, ports={"22/tcp": self.ssh_port}
         )
         return container
 
@@ -129,15 +140,42 @@ class LeanEngine(Engine):
         Sets up SSH access to the Docker container, including generating an SSH key pair if necessary,
         and configuring the container to accept SSH connections using the generated key.
         """
-        if not os.path.exists(self.ssh_key_path):
-            subprocess.run(['ssh-keygen', '-t', 'rsa', '-b', '2048', '-f', self.ssh_key_path, '-N', ''], check=True)
-
-        subprocess.run(['docker', 'exec', self.container.id, 'mkdir', '-p', '/root/.ssh'], check=True)
-        subprocess.run(['docker', 'cp', f'{self.ssh_key_path}.pub', f'{self.container.id}:/root/.ssh/authorized_keys'], check=True)
-        subprocess.run(['docker', 'exec', self.container.id, 'chmod', '600', '/root/.ssh/authorized_keys'], check=True)
-        subprocess.run(['docker', 'exec', self.container.id, 'chown', 'root:root', '/root/.ssh/authorized_keys'], check=True)
+        if not self.ssh_key_path.exists():
+            subprocess.run(
+                ["ssh-keygen", "-t", "rsa", "-b", "2048", "-f", str(self.ssh_key_path), "-N", ""],
+                check=True,
+            )
+
+        subprocess.run(
+            ["docker", "exec", self.container.id, "mkdir", "-p", "/root/.ssh"], check=True
+        )
+        public_key_path = self.ssh_key_path.parent / f"{self.ssh_key_path.name}.pub"
+        subprocess.run(
+            [
+                "docker",
+                "cp",
+                str(public_key_path),
+                f"{self.container.id}:/root/.ssh/authorized_keys",
+            ],
+            check=True,
+        )
+        subprocess.run(
+            ["docker", "exec", self.container.id, "chmod", "600", "/root/.ssh/authorized_keys"],
+            check=True,
+        )
+        subprocess.run(
+            [
+                "docker",
+                "exec",
+                self.container.id,
+                "chown",
+                "root:root",
+                "/root/.ssh/authorized_keys",
+            ],
+            check=True,
+        )
 
-    def forward(self, argument: Any) -> Tuple[List[LeanResult], dict]:
+    def forward(self, argument: Any) -> tuple[list[LeanResult], dict]:
         """
         Executes Lean code provided as a string or as an object property.
 
@@ -149,36 +187,36 @@ class LeanEngine(Engine):
         """
         code: str = argument if isinstance(argument, str) else argument.prop.prepared_input
 
-        rsp: Optional[LeanResult] = None
-        err: Optional[str] = None
-        tmpfile_path: Optional[str] = None
-        metadata: Dict[str, Any] = {}
+        rsp: LeanResult | None = None
+        err: str | None = None
+        tmpfile_path: Path | None = None
+        metadata: dict[str, Any] = {}
         try:
             with tempfile.NamedTemporaryFile(delete=False, suffix=".lean") as tmpfile:
                 tmpfile.write(code.encode())
-                tmpfile_path = tmpfile.name
+                tmpfile_path = Path(tmpfile.name)
 
             output, exec_metadata = self._execute_lean(tmpfile_path)
             metadata.update(exec_metadata)
 
             if output:
-                rsp = LeanResult({'output': output})
+                rsp = LeanResult({"output": output})
             else:
-                metadata['status'] = 'no_output'
+                metadata["status"] = "no_output"
         except Exception as e:
             err = str(e)
-            metadata.update({'status': 'error', 'message': err})
-            print(f"Error during Lean execution: {err}")
+            metadata.update({"status": "error", "message": err})
+            UserMessage(f"Error during Lean execution: {err}")
         finally:
-            if tmpfile_path and os.path.exists(tmpfile_path):
-                os.remove(tmpfile_path)
+            if tmpfile_path and tmpfile_path.exists():
+                tmpfile_path.unlink()
             if self.container:
-                print(f"Killing Docker container '{self.container.id}'...")
+                UserMessage(f"Killing Docker container '{self.container.id}'...")
                 self.container.remove(force=True)
 
         return [rsp] if rsp else [], metadata
 
-    def _execute_lean(self, filepath: str) -> Tuple[str, dict]:
+    def _execute_lean(self, filepath: str) -> tuple[str, dict]:
         """
         Executes a Lean script within the Docker container via SSH.
 
@@ -191,23 +229,30 @@ class LeanEngine(Engine):
         try:
             ssh = paramiko.SSHClient()
             ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
-            ssh.connect(self.ssh_host, port=self.ssh_port, username=self.ssh_user, key_filename=self.ssh_key_path)
+            ssh.connect(
+                self.ssh_host,
+                port=self.ssh_port,
+                username=self.ssh_user,
+                key_filename=str(self.ssh_key_path),
+            )
 
             elan_path: str = "/usr/local/elan/bin/elan"
             lean_path: str = "/usr/local/elan/bin/lean"
 
-            stdin, stdout, stderr = ssh.exec_command(f"{elan_path} default stable && {lean_path} --version")
+            _stdin, stdout, stderr = ssh.exec_command(
+                f"{elan_path} default stable && {lean_path} --version"
+            )
             output: str = stdout.read().decode()
             error: str = stderr.read().decode()
-            print("SSH Command Output:", output)
-            print("SSH Command Error:", error)
+            UserMessage(f"SSH Command Output: {output}")
+            UserMessage(f"SSH Command Error: {error}")
 
             sftp = ssh.open_sftp()
-            remote_path: str = f"/root/{os.path.basename(filepath)}"
+            remote_path: str = f"/root/{Path(filepath).name}"
             sftp.put(filepath, remote_path)
             sftp.close()
 
-            stdin, stdout, stderr = ssh.exec_command(f"{lean_path} {remote_path}")
+            _stdin, stdout, stderr = ssh.exec_command(f"{lean_path} {remote_path}")
             output = stdout.read().decode()
             error = stderr.read().decode()
 
@@ -215,14 +260,13 @@ class LeanEngine(Engine):
             ssh.close()
 
             if "error" in output.lower() or "error" in error.lower():
-                return output, {'status': 'failure'}
-            elif not output and not error:
-                return "Lean program halted successfully with no output.", {'status': 'success'}
-            else:
-                return output, {'status': 'success'}
+                return output, {"status": "failure"}
+            if not output and not error:
+                return "Lean program halted successfully with no output.", {"status": "success"}
+            return output, {"status": "success"}
 
         except Exception as e:
-            raise RuntimeError(f"SSH command execution failed: {str(e)}")
+            UserMessage(f"SSH command execution failed: {e!s}", raise_with=RuntimeError)
 
     def prepare(self, argument: Any) -> None:
         """

symai/backend/engines/neurosymbolic/__init__.py

@@ -1,8 +1,15 @@
-from ...mixin import (ANTHROPIC_CHAT_MODELS, ANTHROPIC_REASONING_MODELS,
-                      DEEPSEEK_CHAT_MODELS, DEEPSEEK_REASONING_MODELS,
-                      GOOGLE_CHAT_MODELS, GOOGLE_REASONING_MODELS,
-                      GROQ_CHAT_MODELS, GROQ_REASONING_MODELS,
-                      OPENAI_CHAT_MODELS, OPENAI_REASONING_MODELS)
+from ...mixin import (
+    ANTHROPIC_CHAT_MODELS,
+    ANTHROPIC_REASONING_MODELS,
+    DEEPSEEK_CHAT_MODELS,
+    DEEPSEEK_REASONING_MODELS,
+    GOOGLE_CHAT_MODELS,
+    GOOGLE_REASONING_MODELS,
+    GROQ_CHAT_MODELS,
+    GROQ_REASONING_MODELS,
+    OPENAI_CHAT_MODELS,
+    OPENAI_REASONING_MODELS,
+)
 from .engine_anthropic_claudeX_chat import ClaudeXChatEngine
 from .engine_anthropic_claudeX_reasoning import ClaudeXReasoningEngine
 from .engine_deepseekX_reasoning import DeepSeekXReasoningEngine
@@ -13,12 +20,33 @@ from .engine_openai_gptX_reasoning import GPTXReasoningEngine
 
 # create the mapping
 ENGINE_MAPPING = {
-    **{model_name: ClaudeXChatEngine for model_name in ANTHROPIC_CHAT_MODELS},
-    **{model_name: ClaudeXReasoningEngine for model_name in ANTHROPIC_REASONING_MODELS},
-    **{model_name: DeepSeekXReasoningEngine for model_name in DEEPSEEK_REASONING_MODELS},
-    **{model_name: GeminiXReasoningEngine for model_name in GOOGLE_REASONING_MODELS},
-    **{model_name: GPTXChatEngine for model_name in OPENAI_CHAT_MODELS},
-    **{model_name: GPTXReasoningEngine for model_name in OPENAI_REASONING_MODELS},
-    **{model_name: GroqEngine for model_name in GROQ_CHAT_MODELS},
-    **{model_name: GroqEngine for model_name in GROQ_REASONING_MODELS},
+    **dict.fromkeys(ANTHROPIC_CHAT_MODELS, ClaudeXChatEngine),
+    **dict.fromkeys(ANTHROPIC_REASONING_MODELS, ClaudeXReasoningEngine),
+    **dict.fromkeys(DEEPSEEK_REASONING_MODELS, DeepSeekXReasoningEngine),
+    **dict.fromkeys(GOOGLE_REASONING_MODELS, GeminiXReasoningEngine),
+    **dict.fromkeys(OPENAI_CHAT_MODELS, GPTXChatEngine),
+    **dict.fromkeys(OPENAI_REASONING_MODELS, GPTXReasoningEngine),
+    **dict.fromkeys(GROQ_CHAT_MODELS, GroqEngine),
+    **dict.fromkeys(GROQ_REASONING_MODELS, GroqEngine),
 }
+
+__all__ = [
+    "ANTHROPIC_CHAT_MODELS",
+    "ANTHROPIC_REASONING_MODELS",
+    "DEEPSEEK_CHAT_MODELS",
+    "DEEPSEEK_REASONING_MODELS",
+    "ENGINE_MAPPING",
+    "GOOGLE_CHAT_MODELS",
+    "GOOGLE_REASONING_MODELS",
+    "GROQ_CHAT_MODELS",
+    "GROQ_REASONING_MODELS",
+    "OPENAI_CHAT_MODELS",
+    "OPENAI_REASONING_MODELS",
+    "ClaudeXChatEngine",
+    "ClaudeXReasoningEngine",
+    "DeepSeekXReasoningEngine",
+    "GPTXChatEngine",
+    "GPTXReasoningEngine",
+    "GeminiXReasoningEngine",
+    "GroqEngine",
+]