fix: address review feedback

Author: RheagalFire

api/config.py

@@ -127,23 +127,23 @@ def load_generator_config():
     generator_config = load_json_config("generator.json")
 
     # Add client classes to each provider
+    default_map = {
+        "google": GoogleGenAIClient,
+        "openai": OpenAIClient,
+        "openrouter": OpenRouterClient,
+        "ollama": OllamaClient,
+        "bedrock": BedrockClient,
+        "azure": AzureAIClient,
+        "dashscope": DashscopeClient,
+        "litellm": LiteLLMClient,
+    }
     if "providers" in generator_config:
         for provider_id, provider_config in generator_config["providers"].items():
             # Try to set client class from client_class
             if provider_config.get("client_class") in CLIENT_CLASSES:
                 provider_config["model_client"] = CLIENT_CLASSES[provider_config["client_class"]]
             # Fall back to default mapping based on provider_id
-            elif provider_id in ["google", "openai", "openrouter", "ollama", "bedrock", "azure", "dashscope", "litellm"]:
-                default_map = {
-                    "google": GoogleGenAIClient,
-                    "openai": OpenAIClient,
-                    "openrouter": OpenRouterClient,
-                    "ollama": OllamaClient,
-                    "bedrock": BedrockClient,
-                    "azure": AzureAIClient,
-                    "dashscope": DashscopeClient,
-                    "litellm": LiteLLMClient,
-                }
+            elif provider_id in default_map:
                 provider_config["model_client"] = default_map[provider_id]
             else:
                 logger.warning(f"Unknown provider or client class: {provider_id}")

api/litellm_client.py

@@ -167,14 +167,18 @@ def convert_inputs_to_api_kwargs(
         return final_model_kwargs
 
     def parse_chat_completion(self, completion) -> GeneratorOutput:
+        import types
+
         try:
-            data = self.chat_completion_parser(completion)
+            is_stream = isinstance(completion, types.GeneratorType) or type(completion).__name__ == "CustomStreamWrapper"
+            parser = handle_streaming_response if is_stream else self.chat_completion_parser
+            data = parser(completion)
         except Exception as e:
             log.error(f"Error parsing the completion: {e}")
             return GeneratorOutput(data=None, error=str(e), raw_response=completion)
 
         try:
-            usage = self.track_completion_usage(completion)
+            usage = None if is_stream else self.track_completion_usage(completion)
             return GeneratorOutput(
                 data=None, error=None, raw_response=data, usage=usage
             )
@@ -207,7 +211,7 @@ def call(self, api_kwargs: Optional[Dict] = None, model_type: ModelType = ModelT
         import litellm
 
         api_kwargs = api_kwargs or {}
-        log.info(f"api_kwargs: {api_kwargs}")
+        log.debug(f"api_kwargs: {api_kwargs}")
 
         extra: Dict[str, Any] = {}
         if self._api_key:
@@ -218,8 +222,6 @@ def call(self, api_kwargs: Optional[Dict] = None, model_type: ModelType = ModelT
         if model_type == ModelType.EMBEDDER:
             return litellm.embedding(drop_params=True, **api_kwargs, **extra)
         elif model_type == ModelType.LLM:
-            if api_kwargs.get("stream", False):
-                self.chat_completion_parser = handle_streaming_response
             return litellm.completion(drop_params=True, **api_kwargs, **extra)
         else:
             raise ValueError(f"model_type {model_type} is not supported")