fix: add type assertion checks to prevent panics

Boris Atamanovskii created 1 month ago

Change summary

model.go                           |  4 +++-
provider_registry.go               |  2 +-
providers/openai/language_model.go | 18 +++++++++++-------
3 files changed, 15 insertions(+), 9 deletions(-)

Detailed changes

model.go

@@ -34,7 +34,9 @@ type ResponseContent []Content
 func (r ResponseContent) Text() string {
 	for _, c := range r {
 		if c.GetType() == ContentTypeText {
-			return c.(TextContent).Text
+			if textContent, ok := AsContentType[TextContent](c); ok {
+				return textContent.Text
+			}
 		}
 	}
 	return ""

provider_registry.go

@@ -38,7 +38,7 @@ func unmarshalProviderData(data []byte) (ProviderOptionsData, error) {
 		return nil, fmt.Errorf("unknown provider data type: %s", pj.Type)
 	}
 
-	unmarshalFn := val.(UnmarshalFunc)
+	unmarshalFn := val.(UnmarshalFunc) //nolint:forcetypeassert // type enforced by RegisterProviderType
 	return unmarshalFn(pj.Data)
 }
 

providers/openai/language_model.go

@@ -648,14 +648,18 @@ func parseAnnotationsFromDelta(delta openai.ChatCompletionChunkChoiceDelta) []op
 			if annotationMap, ok := annotationData.(map[string]any); ok {
 				if annotationType, ok := annotationMap["type"].(string); ok && annotationType == "url_citation" {
 					if urlCitationData, ok := annotationMap["url_citation"].(map[string]any); ok {
-						annotation := openai.ChatCompletionMessageAnnotation{
-							Type: "url_citation",
-							URLCitation: openai.ChatCompletionMessageAnnotationURLCitation{
-								URL:   urlCitationData["url"].(string),
-								Title: urlCitationData["title"].(string),
-							},
+						url, urlOk := urlCitationData["url"].(string)
+						title, titleOk := urlCitationData["title"].(string)
+						if urlOk && titleOk {
+							annotation := openai.ChatCompletionMessageAnnotation{
+								Type: "url_citation",
+								URLCitation: openai.ChatCompletionMessageAnnotationURLCitation{
+									URL:   url,
+									Title: title,
+								},
+							}
+							annotations = append(annotations, annotation)
 						}
-						annotations = append(annotations, annotation)
 					}
 				}
 			}