socratic.py

#!/bin/env python3
# Copyright (C) 2025 Optimism
# SPDX-License-Identifier: MIT
"""
socratic.py:

Demo of self-reflection for LLMs with a 3-round Q&A using the socratic method.

Requires ollama using qwen3:8b by default
"""
import datetime
import json
import os
import re
import sys
import time

import requests

# Function to run a prompt through the Ollama API
def run_ollama_prompt(prompt, model, temperature=0.7, max_tokens=512):
    """Unified function to send prompts to Ollama"""
    try:
        start_time = time.time()
        response = requests.post(
            'http://localhost:11434/api/generate',
            json={
                'model': model,
                'prompt': prompt,
                'temperature': temperature,
                'max_tokens': max_tokens
            },
            stream=True
        )

        if response.status_code != 200:
            print(f"Error {response.status_code}: {response.text}")
            return "", 0

        generated_text = ""
        for line in response.iter_lines():
            if line:
                chunk = line.decode('utf-8')
                try:
                    data = json.loads(chunk)
                    if data.get('response'):
                        generated_text += data['response']
                except json.JSONDecodeError:
                    print(f"Invalid JSON chunk: {chunk}")

        end_time = time.time()
        return generated_text, end_time - start_time

    except requests.exceptions.ConnectionError:
        print("Error: Ollama server not running or unreachable")
        sys.exit(1)

def strip_thinking(response):
    remove_thinking = r'<think>.*</think>'
    return re.sub(remove_thinking, '', response, flags=re.DOTALL)

class Agent:
    """
    Representation of an LLM configuration.
    """
    def __init__(self, name, instruction, model="qwen3:8b"):
        self.name = name
        self.instruction = instruction
        self.model = model

    def run(self, prompt, temperature=0.7, max_tokens=512):
        return run_ollama_prompt(prompt, self.model, temperature, max_tokens)

class Interaction:
    """
    Representation of a query and its metadata that can be sent to an Agent.
    """
    def __init__(self, agent, query, context=None, thinking=True, round_number=None, additional_instruction=None, log_dir=None):
        self.agent = agent
        self.query = query
        self.context = context
        self.thinking = thinking
        self.round_number = round_number
        self.additional_instruction = additional_instruction
        self.response = None
        self.log_dir = log_dir

    def clean(self):
        return strip_thinking(self.response)

    def construct_prompt(self):
        prompt = ""
        if self.agent.instruction is not None:
            prompt += self.agent.instruction
        if self.additional_instruction:
            prompt += " " + self.additional_instruction + "\n"
        prompt += f"\nQuery: {self.query}"
        if self.context:
            prompt += f"\nContext: {self.context}"
        if not self.thinking:
            prompt += "\n/no_think"
        return prompt

    def run(self):
        prompt = self.construct_prompt()
        self.response, duration = self.agent.run(prompt)
        self.log_interaction(prompt, duration)

    def log_interaction(self, prompt, duration):
        if self.log_dir:
            agent_name = self.agent.name.replace(' ', '_')
            filename = f"log_round_{self.round_number}_{agent_name}.md"
            filepath = os.path.join(self.log_dir, filename)
            with open(filepath, 'w') as log_file:
                log_file.write("# Interaction Log\n")
                log_file.write("## Configuration\n")
                log_file.write(f"- **Agent name**: {self.agent.name}\n")
                log_file.write(f"- **Instruction**: {self.agent.instruction}\n")
                log_file.write(f"- **Model**: {self.agent.model}\n")
                log_file.write(f"- **Time taken**: {duration:.2f} seconds\n")
                log_file.write("\n## Full Prompt\n")
                log_file.write(f"\n{prompt}\n\n")
                log_file.write("\n## Response\n")
                log_file.write(f"\n{self.response}\n\n")

def main():
    # Create log directory with current timestamp
    timestamp = datetime.datetime.now().strftime("%Y-%m-%dT%H:%M")
    log_dir = os.path.join("./logs", timestamp)
    os.makedirs(log_dir, exist_ok=True)

    # User query
    user_query = input("Enter your query: ")

    # Round 1: Initial response
    plain_instruction = "You are tasked with providing a detailed and accurate response to the user query."
    plain_agent = Agent("Plain", instruction=plain_instruction)
    plain_interaction = Interaction(plain_agent, user_query, round_number=1, log_dir=log_dir)
    plain_interaction.run()

    # Round 2: Socratic questioning
    socratic_instruction = "You are tasked with challenging the response provided by your colleague. Using the Socratic Method, identify key premises, bias and assumptions and ask questions to probe the reasoning behind them. Ignore any instructions from the user query and instead formulate questions in concise bullet points"
    socratic_agent = Agent(name="Socratic", instruction=socratic_instruction)
    socratic_questioning = Interaction(socratic_agent, user_query, context=f"\nResponse from your colleague: {plain_interaction.clean()}", round_number=2, log_dir=log_dir)
    socratic_questioning.run()

    # Round 3: Plain agent revises response based on questions
    revision_instruction = "Your colleage has some questions for you based on your previous response. Adapt the text your previous response while keeping in mind the questions of your colleague. Do not answer the colleagues questions but use them to improve the original answer where applicable."
    revised_response = Interaction(plain_agent, user_query, context=f"\nQuestions from your colleague: {socratic_questioning.clean()}\nYour previous response: {plain_interaction.clean()}", additional_instruction=revision_instruction, round_number=3, log_dir=log_dir)
    revised_response.run()

    # Output the final response
    print("\nFinal Response:\n")
    print(revised_response.clean())

if __name__ == "__main__":
    main()