| 6 | logger = logging.getLogger(__name__) |
| 7 | |
| 8 | class PlanSearch: |
| 9 | def __init__(self, system_prompt: str, client, model: str, request_config: dict = None, request_id: str = None): |
| 10 | self.system_prompt = system_prompt |
| 11 | self.client = client |
| 12 | self.model = model |
| 13 | self.request_id = request_id |
| 14 | self.plansearch_completion_tokens = 0 |
| 15 | |
| 16 | # Extract max_tokens from request_config with default |
| 17 | self.max_tokens = 4096 |
| 18 | if request_config: |
| 19 | self.max_tokens = request_config.get('max_tokens', self.max_tokens) |
| 20 | |
| 21 | def generate_observations(self, problem: str, num_observations: int = 3) -> List[str]: |
| 22 | prompt = f"""You are an expert Python programmer. You will be given a competitive programming question |
| 23 | (problem specification). You will return several useful, non-obvious, and correct observations |
| 24 | about the problem, like hints to solve the problem. You will NOT return any code. Be as |
| 25 | creative as possible, going beyond what you think is intuitively correct. |
| 26 | |
| 27 | Here is the competitive programming problem: |
| 28 | {problem} |
| 29 | |
| 30 | Please provide {num_observations} observations.""" |
| 31 | |
| 32 | # Prepare request for logging |
| 33 | provider_request = { |
| 34 | "model": self.model, |
| 35 | "max_tokens": self.max_tokens, |
| 36 | "messages": [ |
| 37 | {"role": "system", "content": self.system_prompt}, |
| 38 | {"role": "user", "content": prompt} |
| 39 | ] |
| 40 | } |
| 41 | |
| 42 | response = self.client.chat.completions.create(**provider_request) |
| 43 | |
| 44 | # Log provider call if conversation logging is enabled |
| 45 | if hasattr(optillm, 'conversation_logger') and optillm.conversation_logger and self.request_id: |
| 46 | response_dict = response.model_dump() if hasattr(response, 'model_dump') else response |
| 47 | optillm.conversation_logger.log_provider_call(self.request_id, provider_request, response_dict) |
| 48 | self.plansearch_completion_tokens += response.usage.completion_tokens |
| 49 | |
| 50 | # Check for valid response with None-checking |
| 51 | if (response is None or |
| 52 | not response.choices or |
| 53 | response.choices[0].message.content is None or |
| 54 | response.choices[0].finish_reason == "length"): |
| 55 | logger.warning("Observations response truncated or empty, returning empty list") |
| 56 | return [] |
| 57 | |
| 58 | observations = response.choices[0].message.content.strip().split('\n') |
| 59 | return [obs.strip() for obs in observations if obs.strip()] |
| 60 | |
| 61 | def generate_derived_observations(self, problem: str, observations: List[str], num_new_observations: int = 2) -> List[str]: |
| 62 | prompt = f"""You are an expert Python programmer. You will be given a competitive programming question |
| 63 | (problem specification) and several correct observations about the problem. |
| 64 | You will brainstorm several new, useful, and correct observations about the problem, derived |
| 65 | from the given observations. You will NOT return any code. Be as creative as possible, going |