diff --git a/griptape/drivers/__init__.py b/griptape/drivers/__init__.py index f95aba1cc..64cd16c80 100644 --- a/griptape/drivers/__init__.py +++ b/griptape/drivers/__init__.py @@ -14,6 +14,7 @@ from .prompt.dummy import DummyPromptDriver from .prompt.ollama import OllamaPromptDriver from .prompt.grok import GrokPromptDriver +from .prompt.astraflow import AstraflowPromptDriver from .prompt.griptape_cloud import GriptapeCloudPromptDriver from .prompt.perplexity import PerplexityPromptDriver @@ -144,6 +145,7 @@ "AmazonSageMakerJumpstartEmbeddingDriver", "AmazonSageMakerJumpstartPromptDriver", "AmazonSqsEventListenerDriver", + "AstraflowPromptDriver", "AnthropicPromptDriver", "AstraDbVectorStoreDriver", "AwsIotCoreEventListenerDriver", diff --git a/griptape/drivers/prompt/astraflow/__init__.py b/griptape/drivers/prompt/astraflow/__init__.py new file mode 100644 index 000000000..830ad8110 --- /dev/null +++ b/griptape/drivers/prompt/astraflow/__init__.py @@ -0,0 +1,5 @@ +from griptape.drivers.prompt.astraflow_prompt_driver import AstraflowPromptDriver + +__all__ = [ + "AstraflowPromptDriver", +] diff --git a/griptape/drivers/prompt/astraflow_prompt_driver.py b/griptape/drivers/prompt/astraflow_prompt_driver.py new file mode 100644 index 000000000..c3663b186 --- /dev/null +++ b/griptape/drivers/prompt/astraflow_prompt_driver.py @@ -0,0 +1,17 @@ +from __future__ import annotations + +from attrs import define, field + +from griptape.drivers.prompt.openai import OpenAiChatPromptDriver + + +@define +class AstraflowPromptDriver(OpenAiChatPromptDriver): + """Prompt Driver for Astraflow (UCloud / 优刻得). + + Astraflow is an OpenAI-compatible AI model aggregation platform supporting 200+ models. + Set ``base_url`` to ``https://api.modelverse.cn/v1`` and use ``ASTRAFLOW_CN_API_KEY`` + to use the China endpoint instead of the global endpoint. + """ + + base_url: str = field(default="https://api-us-ca.umodelverse.ai/v1", kw_only=True, metadata={"serializable": True})