phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-26 03:04:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Refactor ASI transformation code to improve maintainability and fix linting issues

Browse source
This commit is contained in:
rajashekarcs2023 2025-04-14 23:06:36 -07:00
parent bd16dc6558
commit 038db90428
1 changed files with 72 additions and 72 deletions
Download patch file Download diff file Expand all files Collapse all files

120
litellm/llms/asi/chat/transformation.py
View file

@ -33,7 +33,7 @@ class ASIChatCompletionStreamingHandler(BaseModelResponseIterator):
model=chunk.get("model", ""), # Use empty string as fallback model=chunk.get("model", ""), # Use empty string as fallback
choices=chunk.get("choices", []), choices=chunk.get("choices", []),
) )
except Exception as e: except Exception:
# Log the error but don't crash - silently continue # Log the error but don't crash - silently continue
# We don't have access to logging_obj here, so we can't log the error # We don't have access to logging_obj here, so we can't log the error
# Return a minimal valid response # Return a minimal valid response
@ -63,6 +63,60 @@ class ASIChatConfig(OpenAIGPTConfig):
"""Get the API base URL for ASI""" """Get the API base URL for ASI"""
return api_base or get_secret_str("ASI_API_BASE") or "https://api.asi1.ai/v1" return api_base or get_secret_str("ASI_API_BASE") or "https://api.asi1.ai/v1"
def _extract_json_from_message_content(self, content: str) -> Optional[str]:
"""Extract JSON from message content if possible"""
if content and isinstance(content, str):
return extract_json(content)
return None
def _process_model_response_choices(self, choices, logging_obj: Any) -> None:
"""Process choices from a ModelResponse object"""
for choice in choices:
try:
# For non-streaming responses (message)
if hasattr(choice, "message"):
message = getattr(choice, "message")
if hasattr(message, "content"):
content = getattr(message, "content")
extracted_json = self._extract_json_from_message_content(content)
if extracted_json:
setattr(message, "content", extracted_json)
# For streaming responses (delta)
elif hasattr(choice, "delta"):
delta = getattr(choice, "delta")
if hasattr(delta, "content"):
content = getattr(delta, "content")
extracted_json = self._extract_json_from_message_content(content)
if extracted_json:
setattr(delta, "content", extracted_json)
if hasattr(logging_obj, "verbose") and logging_obj.verbose:
logging_obj.debug("ASI: Successfully extracted JSON from streaming")
except Exception as attr_error:
# Log attribute access errors but continue processing
if hasattr(logging_obj, "verbose") and logging_obj.verbose:
logging_obj.debug(f"ASI: Error accessing attributes: {str(attr_error)}")
def _process_dict_response_choices(self, choices, logging_obj: Any) -> None:
"""Process choices from a dictionary response"""
for choice in choices:
if not isinstance(choice, dict):
continue
# Handle delta for streaming
if "delta" in choice and isinstance(choice["delta"], dict) and "content" in choice["delta"]:
content = choice["delta"]["content"]
extracted_json = self._extract_json_from_message_content(content)
if extracted_json:
choice["delta"]["content"] = extracted_json
# Handle message for non-streaming
elif "message" in choice and isinstance(choice["message"], dict) and "content" in choice["message"]:
content = choice["message"]["content"]
extracted_json = self._extract_json_from_message_content(content)
if extracted_json:
choice["message"]["content"] = extracted_json
def transform_response( def transform_response(
self, self,
model: str, model: str,
@ -96,75 +150,21 @@ class ASIChatConfig(OpenAIGPTConfig):
# Check if JSON extraction is requested # Check if JSON extraction is requested
json_requested = optional_params.get("json_response_requested", False) or json_mode json_requested = optional_params.get("json_response_requested", False) or json_mode
if json_requested: if not json_requested:
return response
if hasattr(logging_obj, "verbose") and logging_obj.verbose: if hasattr(logging_obj, "verbose") and logging_obj.verbose:
logging_obj.debug("ASI: JSON response format requested, applying extraction") logging_obj.debug("ASI: JSON response format requested, applying extraction")
try: try:
# For ModelResponse objects, directly access the choices # For ModelResponse objects, directly access the choices
if isinstance(response, ModelResponse) and hasattr(response, "choices"): if isinstance(response, ModelResponse) and hasattr(response, "choices"):
choices = response.choices self._process_model_response_choices(response.choices, logging_obj)
# Process each choice
for choice in choices:
# Use type-safe attribute access with proper guards
try:
# Try to extract content from message or delta
content = None
# For non-streaming responses (message)
if hasattr(choice, "message"):
message = getattr(choice, "message")
if hasattr(message, "content"):
content = getattr(message, "content")
# Apply JSON extraction if content exists
if content and isinstance(content, str):
extracted_json = extract_json(content)
if extracted_json:
# Update content safely
setattr(message, "content", extracted_json)
if hasattr(logging_obj, "verbose") and logging_obj.verbose:
print(f"ASI: Successfully extracted JSON: {extracted_json[:100]}...")
# For streaming responses (delta)
elif hasattr(choice, "delta"):
delta = getattr(choice, "delta")
if hasattr(delta, "content"):
content = getattr(delta, "content")
# Apply JSON extraction if content exists
if content and isinstance(content, str):
extracted_json = extract_json(content)
if extracted_json:
# Update content safely
setattr(delta, "content", extracted_json)
if hasattr(logging_obj, "verbose") and logging_obj.verbose:
logging_obj.debug(f"ASI: Successfully extracted JSON from streaming: {extracted_json[:100]}...")
except Exception as attr_error:
# Log attribute access errors but continue processing
if hasattr(logging_obj, "verbose") and logging_obj.verbose:
logging_obj.debug(f"ASI: Error accessing attributes: {str(attr_error)}")
# For streaming responses, handle delta content # For streaming responses, handle delta content
elif isinstance(response, dict) and "choices" in response: elif isinstance(response, dict) and "choices" in response:
for choice in response["choices"]: self._process_dict_response_choices(response["choices"], logging_obj)
if isinstance(choice, dict):
# Handle delta for streaming
if "delta" in choice and isinstance(choice["delta"], dict) and "content" in choice["delta"]:
content = choice["delta"]["content"]
if content:
extracted_json = extract_json(content)
if extracted_json:
choice["delta"]["content"] = extracted_json
# Handle message for non-streaming
elif "message" in choice and isinstance(choice["message"], dict) and "content" in choice["message"]:
content = choice["message"]["content"]
if content:
extracted_json = extract_json(content)
if extracted_json:
choice["message"]["content"] = extracted_json
except Exception as e: except Exception as e:
# Log the error but don't fail the request # Log the error but don't fail the request
if hasattr(logging_obj, "verbose") and logging_obj.verbose: if hasattr(logging_obj, "verbose") and logging_obj.verbose: