Untitled

 avatar
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