feat: refactor model hash lookup to improve error handling and code clarity

- Simplify error handling logic by checking for "not found" message directly
- Extract model data fetching into separate _fetch_model_data method
- Extract version enrichment into separate _enrich_version_with_model_data method
- Improve logging consistency using %s formatting
- Rename variables for better clarity (result -> version, e -> exc)
This commit is contained in:
Will Miao
2025-11-06 22:41:50 +08:00
parent ba6e2eadba
commit 4ed3bd9039

View File

@@ -104,44 +104,32 @@ class CivitaiClient:
async def get_model_by_hash(self, model_hash: str) -> Tuple[Optional[Dict], Optional[str]]: async def get_model_by_hash(self, model_hash: str) -> Tuple[Optional[Dict], Optional[str]]:
try: try:
success, result = await self._make_request( success, version = await self._make_request(
'GET', 'GET',
f"{self.base_url}/model-versions/by-hash/{model_hash}", f"{self.base_url}/model-versions/by-hash/{model_hash}",
use_auth=True use_auth=True
) )
if success: if not success:
# Get model ID from version data message = str(version)
model_id = result.get('modelId') if "not found" in message.lower():
if model_id: return None, "Model not found"
# Fetch additional model metadata
success_model, data = await self._make_request(
'GET',
f"{self.base_url}/models/{model_id}",
use_auth=True
)
if success_model:
# Enrich version_info with model data
result['model']['description'] = data.get("description")
result['model']['tags'] = data.get("tags", [])
# Add creator from model data logger.error("Failed to fetch model info for %s: %s", model_hash[:10], message)
result['creator'] = data.get("creator") return None, message
self._remove_comfy_metadata(result) model_id = version.get('modelId')
return result, None if model_id:
model_data = await self._fetch_model_data(model_id)
# Handle specific error cases if model_data:
if "not found" in str(result): self._enrich_version_with_model_data(version, model_data)
return None, "Model not found"
self._remove_comfy_metadata(version)
# Other error cases return version, None
logger.error(f"Failed to fetch model info for {model_hash[:10]}: {result}")
return None, str(result)
except RateLimitError: except RateLimitError:
raise raise
except Exception as e: except Exception as exc:
logger.error(f"API Error: {str(e)}") logger.error("API Error: %s", exc)
return None, str(e) return None, str(exc)
async def download_preview_image(self, image_url: str, save_path: str): async def download_preview_image(self, image_url: str, save_path: str):
try: try: