From 86baae9c4db2cb1c57bd649852aac5f9d5ac80b3 Mon Sep 17 00:00:00 2001
From: Chester Curme <chester.curme@gmail.com>
Date: Tue, 11 Mar 2025 19:21:36 -0400
Subject: [PATCH] add test for file search

---
 .../chat_models/test_base.py                  | 31 +++++++++++++++++--
 1 file changed, 29 insertions(+), 2 deletions(-)

diff --git a/libs/partners/openai/tests/integration_tests/chat_models/test_base.py b/libs/partners/openai/tests/integration_tests/chat_models/test_base.py
index 28d14b3b89f..89f9e6e35d5 100644
--- a/libs/partners/openai/tests/integration_tests/chat_models/test_base.py
+++ b/libs/partners/openai/tests/integration_tests/chat_models/test_base.py
@@ -2,6 +2,7 @@
 
 import base64
 import json
+import os
 from pathlib import Path
 from textwrap import dedent
 from typing import Any, AsyncIterator, List, Literal, Optional, cast
@@ -1238,8 +1239,17 @@ def _check_response(response: Optional[BaseMessage]) -> None:
         if block["type"] == "text":
             assert isinstance(block["text"], str)
             for annotation in block["annotations"]:
-                for key in ["end_index", "start_index", "title", "type", "url"]:
-                    assert key in annotation
+                if annotation["type"] == "file_citation":
+                    assert all(
+                        key in annotation
+                        for key in ["file_id", "filename", "index", "type"]
+                    )
+                elif annotation["type"] == "web_search":
+                    assert all(
+                        key in annotation
+                        for key in ["end_index", "start_index", "title", "type", "url"]
+                    )
+
     text_content = response.text()
     assert isinstance(text_content, str)
     assert text_content
@@ -1293,3 +1303,20 @@ async def test_web_search_async() -> None:
         full = chunk if full is None else full + chunk
     assert isinstance(full, AIMessageChunk)
     _check_response(full)
+
+
+def test_file_search() -> None:
+    llm = ChatOpenAI(model="gpt-4o")
+    tool = {
+        "type": "file_search",
+        "vector_store_ids": [os.environ["OPENAI_VECTOR_STORE_ID"]],
+    }
+    response = llm.invoke("What is deep research by OpenAI?", tools=[tool])
+    _check_response(response)
+
+    full: Optional[BaseMessageChunk] = None
+    for chunk in llm.stream("What is deep research by OpenAI?", tools=[tool]):
+        assert isinstance(chunk, AIMessageChunk)
+        full = chunk if full is None else full + chunk
+    assert isinstance(full, AIMessageChunk)
+    _check_response(full)