Squash merge fix/lazy into refactor/google-prompt-driver

collindutter · collindutter · commit b04d136fed6d · 2025-04-09T15:09:36.000-07:00
diff --git a/docs/griptape-framework/drivers/src/vector_store_drivers_10.py b/docs/griptape-framework/drivers/src/vector_store_drivers_10.py
@@ -1,5 +1,7 @@
 import os
 
+from qdrant_client.http.models import Distance, VectorParams
+
 from griptape.chunkers import TextChunker
 from griptape.drivers.embedding.openai import OpenAiEmbeddingDriver
 from griptape.drivers.vector.qdrant import QdrantVectorStoreDriver
@@ -27,7 +29,7 @@
 # Recreate Qdrant collection
 vector_store_driver.client.recreate_collection(
     collection_name=vector_store_driver.collection_name,
-    vectors_config={"size": 1536, "distance": vector_store_driver.distance},
+    vectors_config=VectorParams(size=1536, distance=Distance.COSINE),
 )
 
 # Upsert Artifacts into the Vector Store Driver
diff --git a/griptape/drivers/assistant/openai_assistant_driver.py b/griptape/drivers/assistant/openai_assistant_driver.py
@@ -59,19 +59,26 @@ def client(self) -> openai.OpenAI:
         )
 
     def try_run(self, *args: BaseArtifact) -> TextArtifact:
-        if self.thread_id is None and self.auto_create_thread:
-            self.thread_id = self.client.beta.threads.create().id
-        response = self._create_run(*args)
+        if self.thread_id is None:
+            if self.auto_create_thread:
+                thread_id = self.client.beta.threads.create().id
+                self.thread_id = thread_id
+            else:
+                raise ValueError("Thread ID is required but not provided and auto_create_thread is disabled.")
+        else:
+            thread_id = self.thread_id
+
+        response = self._create_run(thread_id, *args)
 
         response.meta.update({"assistant_id": self.assistant_id, "thread_id": self.thread_id})
 
         return response
 
-    def _create_run(self, *args: BaseArtifact) -> TextArtifact:
+    def _create_run(self, thread_id: str, *args: BaseArtifact) -> TextArtifact:
         content = "\n".join(arg.value for arg in args)
-        message_id = self.client.beta.threads.messages.create(thread_id=self.thread_id, role="user", content=content)
+        message_id = self.client.beta.threads.messages.create(thread_id=thread_id, role="user", content=content)
         with self.client.beta.threads.runs.stream(
-            thread_id=self.thread_id,
+            thread_id=thread_id,
             assistant_id=self.assistant_id,
             event_handler=self.event_handler,
         ) as stream:
@@ -80,7 +87,9 @@ def _create_run(self, *args: BaseArtifact) -> TextArtifact:
 
             message_contents = []
             for message in last_messages:
-                message_contents.append("".join(content.text.value for content in message.content))
+                message_contents.append(
+                    "".join(content.text.value for content in message.content if content.type == "TextContentBlock")
+                )
             message_text = "\n".join(message_contents)
 
             response = TextArtifact(message_text)
diff --git a/griptape/drivers/embedding/amazon_bedrock_cohere_embedding_driver.py b/griptape/drivers/embedding/amazon_bedrock_cohere_embedding_driver.py
@@ -12,7 +12,7 @@
 
 if TYPE_CHECKING:
     import boto3
-    from mypy_boto3_bedrock import BedrockClient
+    from mypy_boto3_bedrock_runtime import BedrockRuntimeClient
 
     from griptape.tokenizers.base_tokenizer import BaseTokenizer
 
@@ -40,12 +40,12 @@ class AmazonBedrockCohereEmbeddingDriver(BaseEmbeddingDriver):
         default=Factory(lambda self: AmazonBedrockTokenizer(model=self.model), takes_self=True),
         kw_only=True,
     )
-    _client: Optional[BedrockClient] = field(
+    _client: Optional[BedrockRuntimeClient] = field(
         default=None, kw_only=True, alias="client", metadata={"serializable": False}
     )
 
     @lazy_property()
-    def client(self) -> BedrockClient:
+    def client(self) -> BedrockRuntimeClient:
         return self.session.client("bedrock-runtime")
 
     def try_embed_chunk(self, chunk: str) -> list[float]:
diff --git a/griptape/drivers/embedding/amazon_bedrock_titan_embedding_driver.py b/griptape/drivers/embedding/amazon_bedrock_titan_embedding_driver.py
@@ -14,7 +14,7 @@
 
 if TYPE_CHECKING:
     import boto3
-    from mypy_boto3_bedrock import BedrockClient
+    from mypy_boto3_bedrock_runtime import BedrockRuntimeClient
 
     from griptape.tokenizers.base_tokenizer import BaseTokenizer
 
@@ -38,12 +38,12 @@ class AmazonBedrockTitanEmbeddingDriver(BaseEmbeddingDriver):
         default=Factory(lambda self: AmazonBedrockTokenizer(model=self.model), takes_self=True),
         kw_only=True,
     )
-    _client: Optional[BedrockClient] = field(
+    _client: Optional[BedrockRuntimeClient] = field(
         default=None, kw_only=True, alias="client", metadata={"serializable": False}
     )
 
     @lazy_property()
-    def client(self) -> BedrockClient:
+    def client(self) -> BedrockRuntimeClient:
         return self.session.client("bedrock-runtime")
 
     def try_embed_artifact(self, artifact: TextArtifact | ImageArtifact) -> list[float]:
diff --git a/griptape/drivers/embedding/amazon_sagemaker_jumpstart_embedding_driver.py b/griptape/drivers/embedding/amazon_sagemaker_jumpstart_embedding_driver.py
@@ -11,7 +11,7 @@
 
 if TYPE_CHECKING:
     import boto3
-    from mypy_boto3_sagemaker import SageMakerClient
+    from mypy_boto3_sagemaker_runtime import SageMakerRuntimeClient
 
 
 @define
@@ -20,12 +20,12 @@ class AmazonSageMakerJumpstartEmbeddingDriver(BaseEmbeddingDriver):
     endpoint: str = field(kw_only=True, metadata={"serializable": True})
     custom_attributes: str = field(default="accept_eula=true", kw_only=True, metadata={"serializable": True})
     inference_component_name: Optional[str] = field(default=None, kw_only=True, metadata={"serializable": True})
-    _client: Optional[SageMakerClient] = field(
+    _client: Optional[SageMakerRuntimeClient] = field(
         default=None, kw_only=True, alias="client", metadata={"serializable": False}
     )
 
     @lazy_property()
-    def client(self) -> SageMakerClient:
+    def client(self) -> SageMakerRuntimeClient:
         return self.session.client("sagemaker-runtime")
 
     def try_embed_chunk(self, chunk: str) -> list[float]:
diff --git a/griptape/drivers/embedding/huggingface_hub_embedding_driver.py b/griptape/drivers/embedding/huggingface_hub_embedding_driver.py
@@ -37,4 +37,4 @@ def client(self) -> InferenceClient:
     def try_embed_chunk(self, chunk: str) -> list[float]:
         response = self.client.feature_extraction(chunk)
 
-        return response.flatten().tolist()
+        return [float(val) for val in response.flatten().tolist()]
diff --git a/griptape/drivers/event_listener/amazon_sqs_event_listener_driver.py b/griptape/drivers/event_listener/amazon_sqs_event_listener_driver.py
@@ -10,8 +10,11 @@
 from griptape.utils.decorators import lazy_property
 
 if TYPE_CHECKING:
+    from collections.abc import Sequence
+
     import boto3
     from mypy_boto3_sqs import SQSClient
+    from mypy_boto3_sqs.type_defs import SendMessageBatchRequestEntryTypeDef
 
 
 @define
@@ -28,7 +31,7 @@ def try_publish_event_payload(self, event_payload: dict) -> None:
         self.client.send_message(QueueUrl=self.queue_url, MessageBody=json.dumps(event_payload))
 
     def try_publish_event_payload_batch(self, event_payload_batch: list[dict]) -> None:
-        entries = [
+        entries: Sequence[SendMessageBatchRequestEntryTypeDef] = [
             {"Id": str(event_payload["id"]), "MessageBody": json.dumps(event_payload)}
             for event_payload in event_payload_batch
         ]
diff --git a/griptape/drivers/file_manager/amazon_s3_file_manager_driver.py b/griptape/drivers/file_manager/amazon_s3_file_manager_driver.py
@@ -12,6 +12,7 @@
 if TYPE_CHECKING:
     import boto3
     from mypy_boto3_s3 import S3Client
+    from mypy_boto3_s3.type_defs import PaginatorConfigTypeDef
 
 
 @define
@@ -98,7 +99,7 @@ def _to_dir_full_key(self, path: str) -> str:
 
     def _list_files_and_dirs(self, full_key: str, **kwargs) -> list[str]:
         max_items = kwargs.get("max_items")
-        pagination_config = {}
+        pagination_config: PaginatorConfigTypeDef = {}
         if max_items is not None:
             pagination_config["MaxItems"] = max_items
 
@@ -112,12 +113,12 @@ def _list_files_and_dirs(self, full_key: str, **kwargs) -> list[str]:
         files_and_dirs = []
         for page in pages:
             for obj in page.get("CommonPrefixes", []):
-                prefix = obj.get("Prefix")
+                prefix = obj.get("Prefix", "")
                 directory = prefix[len(full_key) :].rstrip("/")
                 files_and_dirs.append(directory)
 
             for obj in page.get("Contents", []):
-                key = obj.get("Key")
+                key = obj.get("Key", "")
                 file = key[len(full_key) :]
                 files_and_dirs.append(file)
         return files_and_dirs
diff --git a/griptape/drivers/image_generation/amazon_bedrock_image_generation_driver.py b/griptape/drivers/image_generation/amazon_bedrock_image_generation_driver.py
@@ -12,7 +12,7 @@
 
 if TYPE_CHECKING:
     import boto3
-    from mypy_boto3_bedrock import BedrockClient
+    from mypy_boto3_bedrock_runtime import BedrockRuntimeClient
 
 
 @define
@@ -32,12 +32,12 @@ class AmazonBedrockImageGenerationDriver(BaseMultiModelImageGenerationDriver):
     image_width: int = field(default=512, kw_only=True, metadata={"serializable": True})
     image_height: int = field(default=512, kw_only=True, metadata={"serializable": True})
     seed: Optional[int] = field(default=None, kw_only=True, metadata={"serializable": True})
-    _client: Optional[BedrockClient] = field(
+    _client: Optional[BedrockRuntimeClient] = field(
         default=None, kw_only=True, alias="client", metadata={"serializable": False}
     )
 
     @lazy_property()
-    def client(self) -> BedrockClient:
+    def client(self) -> BedrockRuntimeClient:
         return self.session.client("bedrock-runtime")
 
     def try_text_to_image(self, prompts: list[str], negative_prompts: Optional[list[str]] = None) -> ImageArtifact:
diff --git a/griptape/drivers/memory/conversation/amazon_dynamodb_conversation_memory_driver.py b/griptape/drivers/memory/conversation/amazon_dynamodb_conversation_memory_driver.py
@@ -45,7 +45,7 @@ def load(self) -> tuple[list[Run], dict[str, Any]]:
         response = self.table.get_item(Key=self._get_key())
 
         if "Item" in response and self.value_attribute_key in response["Item"]:
-            memory_dict = json.loads(response["Item"][self.value_attribute_key])
+            memory_dict = json.loads(str(response["Item"][self.value_attribute_key]))
             return self._from_params_dict(memory_dict)
         return [], {}
 
diff --git a/griptape/drivers/sql/amazon_redshift_sql_driver.py b/griptape/drivers/sql/amazon_redshift_sql_driver.py
@@ -72,7 +72,7 @@ def execute_query_raw(self, query: str) -> Optional[list[dict[str, Optional[Any]
         if self.database_credentials_secret_arn:
             function_kwargs["SecretArn"] = self.database_credentials_secret_arn
 
-        response = self.client.execute_statement(**function_kwargs)
+        response = self.client.execute_statement(**function_kwargs)  # pyright: ignore[reportArgumentType]
         response_id = response["Id"]
 
         statement = self.client.describe_statement(Id=response_id)
@@ -92,7 +92,7 @@ def execute_query_raw(self, query: str) -> Optional[list[dict[str, Optional[Any]
                 )
                 results = results + response.get("Records", [])
 
-            return self._post_process(statement_result["ColumnMetadata"], results)
+            return self._post_process(statement_result["ColumnMetadata"], results)  # pyright: ignore[reportArgumentType]
 
         if statement["Status"] in ["FAILED", "ABORTED"]:
             return None
@@ -110,5 +110,5 @@ def get_table_schema(self, table_name: str, schema: Optional[str] = None) -> Opt
             function_kwargs["DbUser"] = self.db_user
         if self.database_credentials_secret_arn:
             function_kwargs["SecretArn"] = self.database_credentials_secret_arn
-        response = self.client.describe_table(**function_kwargs)
-        return str([col["name"] for col in response["ColumnList"]])
+        response = self.client.describe_table(**function_kwargs)  # pyright: ignore[reportArgumentType]
+        return str([col["name"] for col in response["ColumnList"] if "name" in col])
diff --git a/griptape/drivers/vector/marqo_vector_store_driver.py b/griptape/drivers/vector/marqo_vector_store_driver.py
@@ -201,7 +201,7 @@ def query(
             "filter_string": f"namespace:{namespace}" if namespace else None,
         } | kwargs
 
-        results = self.client.index(self.index).search(query, **params)
+        results = self.client.index(self.index).search(str(query), **params)
         return self.__process_results(results, include_vectors=include_vectors)
 
     def delete_index(self, name: str) -> dict[str, Any]:
diff --git a/griptape/drivers/vector/pinecone_vector_store_driver.py b/griptape/drivers/vector/pinecone_vector_store_driver.py
@@ -34,7 +34,7 @@ def client(self) -> pinecone.Pinecone:
         )
 
     @lazy_property()
-    def index(self) -> pinecone.Index:
+    def index(self) -> pinecone.data.index.Index:
         return self.client.Index(self.index_name)
 
     def upsert_vector(
@@ -49,12 +49,12 @@ def upsert_vector(
 
         params: dict[str, Any] = {"namespace": namespace} | kwargs
 
-        self.index.upsert(vectors=[(vector_id, vector, meta)], **params)
+        self.index.upsert(vectors=[(vector_id, vector, meta)], **params)  # pyright: ignore[reportArgumentType]
 
         return vector_id
 
     def load_entry(self, vector_id: str, *, namespace: Optional[str] = None) -> Optional[BaseVectorStoreDriver.Entry]:
-        result = self.index.fetch(ids=[vector_id], namespace=namespace).to_dict()
+        result = self.index.fetch(ids=[vector_id], namespace=namespace).to_dict()  # pyright: ignore[reportAttributeAccessIssue]
         vectors = list(result["vectors"].values())
 
         if len(vectors) > 0:
@@ -85,9 +85,9 @@ def load_entries(self, *, namespace: Optional[str] = None) -> list[BaseVectorSto
                 id=r["id"],
                 vector=r["values"],
                 meta=r["metadata"],
-                namespace=results["namespace"],
+                namespace=results["namespace"],  # pyright: ignore[reportIndexIssue]
             )
-            for r in results["matches"]
+            for r in results["matches"]  # pyright: ignore[reportIndexIssue]
         ]
 
     def query_vector(
@@ -115,9 +115,9 @@ def query_vector(
                 vector=r["values"],
                 score=r["score"],
                 meta=r["metadata"],
-                namespace=results["namespace"],
+                namespace=results["namespace"],  # pyright: ignore[reportIndexIssue]
             )
-            for r in results["matches"]
+            for r in results["matches"]  # pyright: ignore[reportIndexIssue]
         ]
 
     def delete_vector(self, vector_id: str) -> NoReturn:
diff --git a/griptape/drivers/vector/qdrant_vector_store_driver.py b/griptape/drivers/vector/qdrant_vector_store_driver.py
@@ -123,12 +123,13 @@ def query_vector(
         # Convert results to QueryResult objects
         return [
             BaseVectorStoreDriver.Entry(
-                id=result.id,
-                vector=result.vector if include_vectors else [],
+                id=str(result.id),
+                vector=result.vector if include_vectors else [],  # pyright: ignore[reportArgumentType]
                 score=result.score,
                 meta={k: v for k, v in result.payload.items() if k not in ["_score", "_tensor_facets"]},
             )
             for result in results
+            if result.payload is not None
         ]
 
     def upsert_vector(
@@ -184,9 +185,11 @@ def load_entry(self, vector_id: str, *, namespace: Optional[str] = None) -> Opti
         results = self.client.retrieve(collection_name=self.collection_name, ids=[vector_id])
         if results:
             entry = results[0]
+            if entry.payload is None:
+                entry.payload = {}
             return BaseVectorStoreDriver.Entry(
-                id=entry.id,
-                vector=entry.vector,
+                id=str(entry.id),
+                vector=entry.vector if entry.vector is not None else [],  # pyright: ignore[reportArgumentType]
                 meta={k: v for k, v in entry.payload.items() if k not in ["_score", "_tensor_facets"]},
             )
         return None
@@ -209,9 +212,10 @@ def load_entries(self, *, namespace: Optional[str] = None, **kwargs) -> list[Bas
 
         return [
             BaseVectorStoreDriver.Entry(
-                id=entry.id,
-                vector=entry.vector if kwargs.get("with_vectors", True) else [],
+                id=str(entry.id),
+                vector=entry.vector if kwargs.get("with_vectors", True) else [],  # pyright: ignore[reportArgumentType]
                 meta={k: v for k, v in entry.payload.items() if k not in ["_score", "_tensor_facets"]},
             )
             for entry in results
+            if entry.payload is not None
         ]
diff --git a/griptape/drivers/vector/redis_vector_store_driver.py b/griptape/drivers/vector/redis_vector_store_driver.py
@@ -85,8 +85,8 @@ def load_entry(self, vector_id: str, *, namespace: Optional[str] = None) -> Opti
         """
         key = self._generate_key(vector_id, namespace)
         result = self.client.hgetall(key)
-        vector = np.frombuffer(result[b"vector"], dtype=np.float32).tolist()
-        meta = json.loads(result[b"metadata"]) if b"metadata" in result else None
+        vector = np.frombuffer(result[b"vector"], dtype=np.float32).tolist()  # pyright: ignore[reportIndexIssue] https://github.com/redis/redis-py/issues/2399
+        meta = json.loads(result[b"metadata"]) if b"metadata" in result else None  # pyright: ignore[reportIndexIssue, reportOperatorIssue]
 
         return BaseVectorStoreDriver.Entry(id=vector_id, meta=meta, vector=vector, namespace=namespace)
 
@@ -100,7 +100,7 @@ def load_entries(self, *, namespace: Optional[str] = None) -> list[BaseVectorSto
         keys = self.client.keys(pattern)
 
         entries = []
-        for key in keys:
+        for key in keys:  # pyright: ignore[reportGeneralTypeIssues] https://github.com/redis/redis-py/issues/2399
             entry = self.load_entry(key.decode("utf-8"), namespace=namespace)
             if entry:
                 entries.append(entry)
@@ -136,7 +136,7 @@ def query_vector(
 
         query_params = {"vector": np.array(vector, dtype=np.float32).tobytes()}
 
-        results = self.client.ft(self.index).search(query_expression, query_params).docs  # pyright: ignore[reportArgumentType]
+        results = self.client.ft(self.index).search(query_expression, query_params).docs  # pyright: ignore[reportArgumentType, reportAttributeAccessIssue]
 
         query_results = []
         for document in results:
diff --git a/griptape/drivers/web_search/exa_web_search_driver.py b/griptape/drivers/web_search/exa_web_search_driver.py
@@ -26,8 +26,8 @@ def client(self) -> Exa:
         return import_optional_dependency("exa_py").Exa(api_key=self.api_key)
 
     def search(self, query: str, **kwargs) -> ListArtifact[JsonArtifact]:
-        response = self.client.search_and_contents(
-            highlights=self.highlights,
+        response = self.client.search_and_contents(  # pyright: ignore[reportCallIssue]
+            highlights=self.highlights,  # pyright: ignore[reportArgumentType]
             use_autoprompt=self.use_autoprompt,
             query=query,
             num_results=self.results_count,
diff --git a/griptape/schemas/base_schema.py b/griptape/schemas/base_schema.py
diff --git a/griptape/utils/decorators.py b/griptape/utils/decorators.py
diff --git a/pyproject.toml b/pyproject.toml
diff --git a/tests/unit/drivers/assistant/test_openai_assistant_driver.py b/tests/unit/drivers/assistant/test_openai_assistant_driver.py
diff --git a/uv.lock b/uv.lock
