handle the server-overloaded case (#261)

mtblanton · web-flow · commit 4eb0dd1a6702 · 2025-04-02T14:31:30.000-05:00
Depends on #260 This PR handles the `server-overloaded` case from the infini-gram API and raises a relevant error for the client.
diff --git a/src/attribution/attribution_service.py b/src/attribution/attribution_service.py
@@ -213,6 +213,12 @@ def get_attribution(
             attribution_response.title,
             str(attribution_response.detail),
         )
+
+        if attribution_response.type_ == "server-overloaded":
+            raise exceptions.ServiceUnavailable(
+                description="OLMoTrace is currently overloaded. Please try again later."
+            )
+
         raise exceptions.InternalServerError(
             description=f"infini-gram API reported an error: {attribution_response.title}"
         )

Original file line number	Diff line number	Diff line change
`@@ -213,6 +213,12 @@ def get_attribution(`
`213`	`213`	`attribution_response.title,`
`214`	`214`	`str(attribution_response.detail),`
`215`	`215`	`)`
	`216`	`+`
	`217`	`+ if attribution_response.type_ == "server-overloaded":`
	`218`	`+ raise exceptions.ServiceUnavailable(`
	`219`	`+ description="OLMoTrace is currently overloaded. Please try again later."`
	`220`	`+ )`
	`221`	`+`
`216`	`222`	`raise exceptions.InternalServerError(`
`217`	`223`	`description=f"infini-gram API reported an error: {attribution_response.title}"`
`218`	`224`	`)`