Update inference.py

inference.py

@@ -1,7 +1,8 @@
 # inference.py
-
-from typing import List, Dict, Generator, Optional
+from typing import List, Dict, Optional
 from hf_client import get_inference_client
+from models import find_model
+
 
 def chat_completion(
     model_id: str,
@@ -14,20 +15,25 @@ def chat_completion(
 
     Args:
         model_id: The model identifier to use.
-        messages: A list of OpenAI‑style {'role':'...', 'content':'...'} messages.
+        messages: A list of OpenAI-style {'role','content'} messages.
         provider: Optional override for provider; uses model default if None.
         max_tokens: Maximum tokens to generate.
 
     Returns:
         The assistant's response content.
     """
-    client = get_inference_client(model_id, provider or "auto")
-    response = client.chat.completions.create(
+    # resolve default provider from registry if needed
+    if provider is None:
+        meta = find_model(model_id)
+        provider = meta.default_provider if meta else "auto"
+
+    client = get_inference_client(model_id, provider)
+    resp = client.chat.completions.create(
         model=model_id,
         messages=messages,
         max_tokens=max_tokens
     )
-    return response.choices[0].message.content
+    return resp.choices[0].message.content
 
 
 def stream_chat_completion(
@@ -35,12 +41,16 @@ def stream_chat_completion(
     messages: List[Dict[str, str]],
     provider: Optional[str] = None,
     max_tokens: int = 4096
-) -> Generator[str, None, None]:
+):
     """
     Generator for streaming chat completions.
     Yields partial message chunks as strings.
     """
-    client = get_inference_client(model_id, provider or "auto")
+    if provider is None:
+        meta = find_model(model_id)
+        provider = meta.default_provider if meta else "auto"
+
+    client = get_inference_client(model_id, provider)
     stream = client.chat.completions.create(
         model=model_id,
         messages=messages,