Untitled
unknown
python
2 months ago
4.0 kB
6
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