Fix/json output (#40)

* added json response mime type to generate_contextual_clarification.

* added structured JSON output support for vLLM client.

* added structured JSON output to get_top_personas and handled dict-wrapped responses.

multi_llm_chatbot_backend/app/core/improved_orchestrator.py

@@ -226,6 +226,7 @@ class ImprovedChatOrchestrator:
                 context=[{"role": "user", "content": user_prompt}],
                 temperature=0.4,
                 max_tokens=1024,
+                response_mime_type="application/json"
             )
 
             cleaned = raw.strip()
@@ -793,7 +794,8 @@ When analyzing the document context:
                 system_prompt=f"You are an assistant that selects the best advisors for a user of {app_title}.",
                 context=[{"role": "user", "content": prompt}],
                 temperature=0.4,
-                max_tokens=150
+                max_tokens=150,
+                response_mime_type="application/json"
             )
 
             # Step 1: Try direct JSON load
@@ -804,6 +806,10 @@ When analyzing the document context:
                 top_ids = re.findall(r'"(.*?)"', llm_response)
                 logger.warning(f"Fallback JSON extraction used: {top_ids}")
 
+            # Handle models that wrap the list in an object (e.g. {"advisor_ids": [...]})
+            if isinstance(top_ids, dict):
+                top_ids = next(iter(top_ids.values()), [])
+
             # Step 3: Filter valid persona IDs
             valid_ids = [pid for pid in top_ids if pid in self.personas]
 

multi_llm_chatbot_backend/app/llm/improved_vllm_client.py

@@ -42,13 +42,18 @@ class ImprovedVllmClient(LLMClient):
             if not self.model_name:
                 await self.refresh_model()
 
-            response = await self.client.chat.completions.create(
+            create_kwargs = dict(
                 model=self.model_name,
                 messages=context_window.messages,
                 temperature=temperature,
                 max_tokens=max_tokens,
             )
 
+            if response_mime_type == "application/json":
+                create_kwargs["response_format"] = {"type": "json_object"}
+
+            response = await self.client.chat.completions.create(**create_kwargs)
+
             text = response.choices[0].message.content.strip()
             return self._clean_response(text)
 

multi_llm_chatbot_backend/app/tests/unit/test_vllm_client.py

@@ -160,6 +160,47 @@ class TestImprovedVllmClient(unittest.TestCase):
         ))
         self.assertIsNone(client.model_name)
 
+    # ------------------------------------------------------------------
+    # generate – response_format for JSON
+    # ------------------------------------------------------------------
+
+    def test_generate_passes_response_format_for_json(self, MockAsyncOpenAI, mock_get_ctx):
+        client = ImprovedVllmClient(
+            api_url=FAKE_URL, api_key=FAKE_KEY, model_name="test-model",
+        )
+        client.client.chat.completions.create = AsyncMock(
+            return_value=_make_completion_mock('{"key": "value"}')
+        )
+
+        asyncio.run(client.generate(
+            system_prompt="Return JSON",
+            context=[{"role": "user", "content": "Hi"}],
+            temperature=0.3,
+            max_tokens=100,
+            response_mime_type="application/json",
+        ))
+
+        call_kwargs = client.client.chat.completions.create.call_args.kwargs
+        self.assertEqual(call_kwargs["response_format"], {"type": "json_object"})
+
+    def test_generate_omits_response_format_when_no_mime_type(self, MockAsyncOpenAI, mock_get_ctx):
+        client = ImprovedVllmClient(
+            api_url=FAKE_URL, api_key=FAKE_KEY, model_name="test-model",
+        )
+        client.client.chat.completions.create = AsyncMock(
+            return_value=_make_completion_mock("plain text response")
+        )
+
+        asyncio.run(client.generate(
+            system_prompt="You are helpful.",
+            context=[{"role": "user", "content": "Hello"}],
+            temperature=0.7,
+            max_tokens=100,
+        ))
+
+        call_kwargs = client.client.chat.completions.create.call_args.kwargs
+        self.assertNotIn("response_format", call_kwargs)
+
     # ------------------------------------------------------------------
     # _clean_response
     # ------------------------------------------------------------------