Renamed mistral tool to mistral_api

tools/mistral_api_infer/pyproject.toml

@@ -0,0 +1,18 @@
+[build-system]
+requires = ["setuptools>=42", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "mistral_api_infer"
+version = "0.1.0"
+requires-python = ">=3.8"
+
+dependencies = [
+    "llm_engine_utils @ file:///root/sia/lib/llm_engine_utils",
+    "mistral-common>=1.0.0",
+    "mistralai>=0.0.7",
+    "python-dotenv>=1.0.0",
+]
+
+[project.scripts]
+mistral_api_infer = "mistral_api_infer.__main__:main"

tools/mistral_api_infer/src/mistral_api_infer/__main__.py

@@ -0,0 +1,42 @@
+from .mistral_llm_engine import MistralApiLlmEngine
+from dotenv import load_dotenv
+from llm_engine_utils.protocol import process
+import argparse
+import os
+
+def main():
+    load_dotenv()
+    parser = argparse.ArgumentParser(description='Mistral API Inference')
+    
+    parser.add_argument(
+        '--model',
+        type=str,
+        default=os.getenv('SIA_MISTRAL_MODEL', 'mistral-large-latest'),
+        help='Model name (default: mistral-large-latest, env: SIA_MISTRAL_MODEL)'
+    )
+
+    parser.add_argument(
+        '--token-limit',
+        type=int,
+        default=int(os.getenv('SIA_MISTRAL_TOKEN_LIMIT', 128000)),
+        help='Token limit for the model (default: 128000, env: SIA_MISTRAL_TOKEN_LIMIT)'
+    )
+
+    parser.add_argument(
+        '--api-key',
+        type=str,
+        default=os.getenv('SIA_MISTRAL_API_KEY'),
+        help='API key for the model (required, env: SIA_MISTRAL_API_KEY)'
+    )
+
+    args = parser.parse_args()
+
+    mistral_llm_engine = MistralLlmEngine(
+        model=args.model,
+        token_limit=args.token_limit,
+        api_key=args.api_key,
+    )
+    process(mistral_llm_engine)
+
+if __name__ == "__main__":
+    main()

tools/mistral_api_infer/src/mistral_api_infer/mistral_llm_engine.py

@@ -0,0 +1,78 @@
+from llm_engine_utils import LlmEngine
+from llm_engine_utils.iterators import skip_prefix
+from mistral_common.protocol.instruct.messages import SystemMessage, UserMessage
+from mistral_common.protocol.instruct.request import ChatCompletionRequest
+from mistral_common.tokens.tokenizers.mistral import MistralTokenizer
+from mistralai import Mistral
+from pathlib import Path
+from typing import Iterator
+
+class MistralLlmEngine(LlmEngine):
+    def __init__(
+            self,
+            model: str,
+            api_key: str,
+            token_limit: int,
+        ):
+        self._model = model
+        self._api_key = api_key
+        self._token_limit = token_limit
+        self._client = Mistral(api_key=api_key)
+        self._tokenizer = MistralTokenizer.v3()
+
+    def infer_xml(self, schema: Path, system: str, context: str, prefix: str) -> Iterator[str]:
+        messages = [
+            {
+                "role": "system",
+                "content": system,
+            },
+            {
+                "role": "user",
+                "content": context,
+            },
+            {
+                "role": "assistant",
+                "content": prefix,
+                "prefix": True,
+            },
+        ] if prefix else [
+            {
+                "role": "system",
+                "content": system,
+            },
+            {
+                "role": "user",
+                "content": context,
+            },
+        ]
+        stream_response = self._client.chat.stream(
+            model=self._model,
+            messages=messages,
+            #temperature=self._temperature,
+        )
+        
+        try:
+            def content_generator():
+                for chunk in stream_response:
+                    if content := chunk.data.choices[0].delta.content:
+                        yield content
+            yield from skip_prefix(content_generator(), prefix)
+        finally:
+            stream_response.response.close()
+
+    def token_count(self, system: str, context: str) -> int:
+        messages = [
+            SystemMessage(content=system),
+            UserMessage(content=context),
+        ]
+
+        tokenized = self._tokenizer.encode_chat_completion(
+            ChatCompletionRequest(
+                messages=messages,
+                model=self._model
+            )
+        )
+        return len(tokenized.tokens)
+
+    def token_limit(self) -> int:
+        return self._token_limit