40 lines
1.3 KiB
Python
40 lines
1.3 KiB
Python
|
|
|
|
from toolserve.sdk import Param, tool, get_secret
|
|
from toolserve.sdk.dataframe import get_df
|
|
import pandas as pd
|
|
import openai
|
|
|
|
@tool
|
|
async def summarize(
|
|
#text: Param(str, "Text to summarize"),
|
|
data_id: Param(int, "ID of the data to summarize"),
|
|
system_prompt: Param(str, "System prompt to use") = "Summarize the following text",
|
|
max_tokens: Param(int, "Maximum number of tokens to generate") = 1000,
|
|
) -> Param(str, "Summarized text"):
|
|
"""Summarize a piece of text using OpenAI Language models.
|
|
|
|
Args:
|
|
text (str): The text to summarize.
|
|
max_tokens (int): The maximum number of tokens to generate.
|
|
|
|
Returns:
|
|
str: The summarized text.
|
|
"""
|
|
df = await get_df(data_id)
|
|
text = df.to_json(orient='records')
|
|
api_key = get_secret("openai_api_key", None)
|
|
model = get_secret("openai_model_summarize", "gpt-4-turbo")
|
|
# Call the OpenAI model with the tools and messages
|
|
messages = [
|
|
{"role": "system", "content": system_prompt},
|
|
{"role": "user", "content": text},
|
|
]
|
|
|
|
client = openai.Client(api_key=api_key)
|
|
completion = openai.chat.completions.create(
|
|
model=model,
|
|
messages=messages,
|
|
)
|
|
summary = completion.choices[0].message.content
|
|
return summary
|