nathadfield commented on code in PR #38736:
URL: https://github.com/apache/airflow/pull/38736#discussion_r1565324842
##########
airflow/providers/openai/hooks/openai.py:
##########
@@ -77,6 +89,165 @@ def get_conn(self) -> OpenAI:
**openai_client_kwargs,
)
+ def create_chat_completion(
+ self,
+ messages: list[
+ ChatCompletionSystemMessageParam
+ | ChatCompletionUserMessageParam
+ | ChatCompletionAssistantMessageParam
+ | ChatCompletionToolMessageParam
+ | ChatCompletionFunctionMessageParam
+ ],
+ model: str = "gpt-3.5-turbo",
+ **kwargs: Any,
+ ) -> list[ChatCompletionMessage]:
+ """
+ Create a model response for the given chat conversation and returns a
list of chat completions.
+
+ :param messages: A list of messages comprising the conversation so far
+ :param model: ID of the model to use
+ """
+ response = self.conn.chat.completions.create(model=model,
messages=messages, **kwargs)
+ return response.choices
+
+ def create_assistant(self, model: str = "gpt-3.5-turbo", **kwargs: Any) ->
Assistant:
+ """Create an OpenAI assistant using the given model.
+
+ :param model: The OpenAI model for the assistant to use.
+ """
+ assistant = self.conn.beta.assistants.create(model=model, **kwargs)
+ return assistant
+
+ def get_assistant(self, assistant_id: str) -> Assistant:
+ """
+ Get an OpenAI assistant.
+
+ :param assistant_id: The ID of the assistant to retrieve.
+ """
+ assistant =
self.conn.beta.assistants.retrieve(assistant_id=assistant_id)
+ return assistant
+
+ def get_assistants(self, **kwargs: Any) -> list[Assistant]:
+ """Get a list of Assistant objects."""
+ assistants = self.conn.beta.assistants.list(**kwargs)
+ return assistants.data
+
+ def get_assistant_by_name(self, assistant_name: str) -> Assistant | None:
+ """Get an OpenAI Assistant object for a given name.
+
+ :param assistant_name: The name of the assistant to retrieve
+ """
+ response = self.get_assistants()
+ for assistant in response:
+ if assistant.name == assistant_name:
+ return assistant
+ return None
+
+ def modify_assistant(self, assistant_id: str, **kwargs: Any) -> Assistant:
+ """Modify an existing Assistant object."""
+ assistant =
self.conn.beta.assistants.update(assistant_id=assistant_id, **kwargs)
+ return assistant
+
+ def delete_assistant(self, assistant_id: str) -> AssistantDeleted:
+ """Delete an OpenAI Assistant for a given ID.
+
+ :param assistant_id: The ID of the assistant to delete.
+ """
+ response = self.conn.beta.assistants.delete(assistant_id=assistant_id)
+ return response
+
+ def create_thread(self, **kwargs: Any) -> Thread:
+ """Create an OpenAI thread."""
+ thread = self.conn.beta.threads.create(**kwargs)
+ return thread
+
+ def modify_thread(self, thread_id: str, metadata: dict) -> Thread:
Review Comment:
Yes, I think that makes sense.
The `metadata` is just a set of K:V pairs that can be attached to the
object. e.g.
```
my_updated_thread = client.beta.threads.update(
"thread_abc123",
metadata={
"modified": "true",
"user": "abc123"
}
)
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscr...@airflow.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
