Untitled
unknown
python
10 months ago
4.0 kB
9
Indexable
async def generate(
self, prompt: str, model: Optional[str] = None, **kwargs: Any
) -> Dict[str, Any]:
"""Generate text using OpenAI API.
Args:
prompt: Input prompt for generation
model: Optional model identifier
**kwargs: Additional parameters for the model
Returns:
Dict containing generation results
"""
try:
messages = []
# Parse system prompt if provided
if "system_message" in kwargs and kwargs["system_message"]:
system_text = str(kwargs["system_message"])
# Remove special tokens if present
system_text = (
system_text.replace("<|im_start|>system\n", "")
.replace("<|im_end|>", "")
.strip()
)
if system_text: # Only add if not empty
messages.append({"role": "system", "content": system_text})
# Parse prompt for history and current message
if prompt:
parts = str(prompt).split("<|im_end|>\n")
for part in parts:
if not part or not part.strip():
continue
if "<|im_start|>user\n" in part:
user_text = part.replace("<|im_start|>user\n", "").strip()
if user_text: # Only add if not empty
messages.append({"role": "user", "content": user_text})
elif "<|im_start|>assistant\n" in part:
assistant_text = part.replace("<|im_start|>assistant\n", "").strip()
if assistant_text: # Only add if not empty
messages.append({"role": "assistant", "content": assistant_text})
# If no special tokens found, treat entire prompt as user message
if not messages or (len(messages) == 1 and messages[0]["role"] == "system"):
messages.append({"role": "user", "content": str(prompt).strip()})
# Validate that we have at least one non-empty message
if not messages:
logger.error("No valid messages found in the prompt")
return {"success": False, "error": "No valid messages found in the prompt"}
logger.info(f"Sending message to OpenAI: {messages}")
# Use Gemini model if not specified
model = model or "gemini-1.5-flash"
if kwargs.get("extract_profile", False):
# Use parse method for structured output
completion = self.client.beta.chat.completions.parse(
model=model,
messages=messages,
response_format=UserProfile,
temperature=kwargs.get("temperature", self.settings.OPENAI_TEMPERATURE),
max_tokens=kwargs.get("max_tokens", self.settings.OPENAI_MAX_TOKENS),
)
response_data = completion.choices[0].message.parsed
logger.info(f"Structured response: {response_data}")
return {"success": True, "data": response_data}
else:
# Use standard completion for text output
response = self.client.chat.completions.create(
model=model,
messages=messages,
temperature=kwargs.get("temperature", self.settings.OPENAI_TEMPERATURE),
max_tokens=kwargs.get("max_tokens", self.settings.OPENAI_MAX_TOKENS),
)
response_text = response.choices[0].message.content
logger.info(f"Response: {response_text}")
return {"success": True, "text": response_text}
except Exception as e:
logger.error(f"Error generating text: {str(e)}")
return {"success": False, "error": str(e)}Editor is loading...
Leave a Comment