Main plugin execution function.
(system_prompt: str, initial_query: str, client, model: str, request_config: dict = None)
| 134 | return None |
| 135 | |
| 136 | def run(system_prompt: str, initial_query: str, client, model: str, request_config: dict = None) -> Tuple[str, int]: |
| 137 | """Main plugin execution function.""" |
| 138 | logger.info("Starting JSON plugin execution") |
| 139 | completion_tokens = 0 |
| 140 | |
| 141 | try: |
| 142 | # Extract schema from response_format in request_config |
| 143 | response_format = request_config.get("response_format") if request_config else None |
| 144 | schema = extract_schema_from_response_format(response_format) |
| 145 | |
| 146 | if not schema: |
| 147 | logger.warning("No valid schema found in response_format") |
| 148 | # Fall back to regular completion if no schema is specified |
| 149 | response = client.chat.completions.create( |
| 150 | model=model, |
| 151 | messages=[ |
| 152 | {"role": "system", "content": system_prompt}, |
| 153 | {"role": "user", "content": initial_query} |
| 154 | ] |
| 155 | ) |
| 156 | return response.choices[0].message.content, response.usage.completion_tokens |
| 157 | |
| 158 | # Initialize JSON generator |
| 159 | json_generator = JSONGenerator() |
| 160 | |
| 161 | # Generate JSON response |
| 162 | result = json_generator.generate_json(initial_query, schema) |
| 163 | |
| 164 | # Convert result to string if it's not already |
| 165 | json_response = json.dumps(result) if isinstance(result, dict) else str(result) |
| 166 | completion_tokens = json_generator.count_tokens(json_response) |
| 167 | |
| 168 | logger.info(f"Successfully generated JSON response: {json_response}") |
| 169 | return json_response, completion_tokens |
| 170 | |
| 171 | except Exception as e: |
| 172 | logger.error(f"Error in JSON plugin: {str(e)}") |
| 173 | # Fall back to regular completion on error |
| 174 | response = client.chat.completions.create( |
| 175 | model=model, |
| 176 | messages=[ |
| 177 | {"role": "system", "content": system_prompt}, |
| 178 | {"role": "user", "content": initial_query} |
| 179 | ] |
| 180 | ) |
| 181 | return response.choices[0].message.content, response.usage.completion_tokens |
nothing calls this directly
no test coverage detected