Update app.py

app.py

@@ -33,14 +33,12 @@ MODELS = {
         },
     },
     "Mistral Small": {
+        # Update these two fields with your Azure values:
         "api_url": "https://ezofisai.services.ai.azure.com/api/projects/firstProject",
-        "model": "mistral-small-2503",
+        "model": "mistral-small-2503",  # this is not used by Azure, just for completeness
         "key_env": "AZUREMIST_API_KEY",
         "response_format": {"type": "json_object"},
-        "extra_headers": {
-            "HTTP-Referer": "https://huggingface.co",
-            "X-Title": "Invoice Extractor",
-        },
+        # No extra_headers needed for Azure
     },
 }
 
@@ -54,17 +52,23 @@ def get_api_key(model_choice):
 def query_llm(model_choice, prompt):
     cfg = MODELS[model_choice]
     headers = {
-        "Authorization": f"Bearer {get_api_key(model_choice)}",
         "Content-Type": "application/json",
     }
+    # Azure OpenAI (Mistral Small) needs api-key header instead of Authorization
+    if model_choice == "mistral-small-2503":
+        headers["api-key"] = get_api_key(model_choice)
+    else:
+        headers["Authorization"] = f"Bearer {get_api_key(model_choice)}"
     if cfg.get("extra_headers"):
         headers.update(cfg["extra_headers"])
     payload = {
-        "model": cfg["model"],
         "messages": [{"role": "user", "content": prompt}],
         "temperature": 0.1,
         "max_tokens": 2000,
     }
+    # Only non-Azure APIs need "model" in payload
+    if model_choice != "mistral-small-2503":
+        payload["model"] = cfg["model"]
     if cfg.get("response_format"):
         payload["response_format"] = cfg["response_format"]
     try:
@@ -206,4 +210,4 @@ with tab2:
     if "last_api" in st.session_state:
         with st.expander("Debug"):
             st.code(st.session_state.last_api)
-            st.code(st.session_state.last_raw)
+            st.code(st.session_state.last_raw)