More pyright fixes

.github/workflows/ci.yml

@@ -39,5 +39,5 @@ jobs:
       - name: Run tests
         run: uv run pytest -v
 
-      # - name: pyright
-      #   run: uv run pyright nilai-api
+      - name: pyright
+        run: uv run pyright

nilai-models/src/nilai_models/daemon.py

@@ -19,6 +19,7 @@ async def get_metadata(num_retries=30):
     service and return as ModelMetadata object"""
     current_retries = 0
     while True:
+        url = None
         try:
             url = f"http://{SETTINGS['host']}:{SETTINGS['port']}/v1/models"
             # Request model metadata from localhost:8000/v1/models
@@ -40,7 +41,10 @@ async def get_metadata(num_retries=30):
                 )
 
         except Exception as e:
-            logger.warning(f"Failed to fetch model metadata from {url}: {e}")
+            if not url:
+                logger.warning(f"Failed to build url: {e}")
+            else:
+                logger.warning(f"Failed to fetch model metadata from {url}: {e}")
             current_retries += 1
             if current_retries >= num_retries:
                 raise e

nilai/__init__.py

@@ -1,4 +0,0 @@
-from nilai_api.app import app as nilai_api_app
-from nilai_models.models.llama_1b_cpu.llama_1b_cpu import app as llama_1b_cpu_app
-
-__all__ = ["nilai_api_app", "llama_1b_cpu_app"]

packages/nilai-common/src/nilai_common/__init__.py

@@ -7,8 +7,8 @@
     Message,
     ModelEndpoint,
     ModelMetadata,
-    CompletionUsage as Usage,
 )
+from openai.types.completion_usage import CompletionUsage as Usage
 from nilai_common.config import SETTINGS
 from nilai_common.discovery import ModelServiceDiscovery
 
@@ -17,9 +17,6 @@
     "ChatRequest",
     "SignedChatCompletion",
     "Choice",
-    "ChoiceChunk",
-    "ChoiceChunkContent",
-    "ChatCompletionChunk",
     "ModelMetadata",
     "Usage",
     "AttestationResponse",

packages/nilai-common/src/nilai_common/api_model.py

@@ -3,7 +3,7 @@
 
 from openai.types.chat import ChatCompletion, ChatCompletionMessage
 from openai.types.chat.chat_completion import Choice as OpenaAIChoice
-from openai.types.chat.chat_completion import CompletionUsage
+from openai.types.completion_usage import CompletionUsage
 from openai.types.chat import ChatCompletionToolParam
 from pydantic import BaseModel, Field
 
@@ -22,7 +22,7 @@
 
 
 class Message(ChatCompletionMessage):
-    role: Literal["system", "user", "assistant"]
+    role: Literal["system", "user", "assistant"]  # type: ignore
 
 
 class Choice(OpenaAIChoice):

packages/nilai-common/src/nilai_common/discovery.py

@@ -171,6 +171,7 @@ async def main():
         license="MIT",
         source="https://github.com/example/model",
         supported_features=["image_classification", "transfer_learning"],
+        tool_support=False,
     )
 
     model_endpoint = ModelEndpoint(

pyproject.toml

@@ -48,4 +48,4 @@ verifier = { workspace = true }
 
 
 [tool.pyright]
-exclude = [".venv"]
+exclude = [".venv", "packages/verifier/"]

tests/__init__.py

@@ -1,3 +1,5 @@
+from openai.types.chat.chat_completion import ChoiceLogprobs
+
 from nilai_common import (
     SignedChatCompletion,
     Message,
@@ -33,9 +35,9 @@
             index=0,
             message=Message(role="assistant", content="test-content"),
             finish_reason="stop",
-            logprobs={"test-logprobs": "test-value"},
+            logprobs=ChoiceLogprobs(),
         )
-    ],
+    ],  # type: ignore
     usage=Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150),
     signature="test-signature",
 )

tests/functional_tests/openai.py

@@ -13,8 +13,10 @@ def test_stream():
     ]
     print(messages)
     response = client.chat.completions.create(
-        model="meta-llama/Llama-3.1-8B-Instruct", messages=messages, stream=True
-    )
+        model="meta-llama/Llama-3.1-8B-Instruct",
+        messages=messages,  # type: ignore
+        stream=True,
+    )  # type: ignore
 
     content = ""
     for chunk in response:
@@ -42,8 +44,10 @@ def test_non_stream():
     ]
     print(messages)
     response = client.chat.completions.create(
-        model="meta-llama/Llama-3.1-8B-Instruct", messages=messages, stream=False
-    )
+        model="meta-llama/Llama-3.1-8B-Instruct",
+        messages=messages,  # type: ignore
+        stream=False,
+    )  # type: ignore
     print(response)
     # raise Exception(f"Response: {response}")
 

tests/functional_tests/tools.py

@@ -35,7 +35,7 @@ def test_tools():
         messages=[
             {"role": "user", "content": "What is the weather like in Paris today?"}
         ],
-        tools=tools,
+        tools=tools,  # type: ignore
     )
     print(response)
     # raise Exception(f"Response: {response}")