import ollama
import ast
import re

def extract_list(input_string):
    # Find the list using regex
    match = re.search(r"\[\[.*?\]\]", input_string, re.DOTALL)
    if match:
        extracted = match.group(0)
        try:
            return ast.literal_eval(extracted)[0]  # Convert string to list
        except (SyntaxError, ValueError):
            return "Error: Could not parse the list."
    return "Error: No list found."

def generate_topic_response(topic):
    query = f'''Generate the 5 topics on the given value: {topic} 
    give the output in the form of a list seperated by commas. in given format
    [["topic1", "topic2", "topic3", "topic4", "topic5"]]
    '''
    # Call the Ollama model without streaming
    response = ollama.chat(
        model='llama3:8b',
        messages=[{'role': 'user', 'content': query}],
        stream=False,  # Set stream to False to get the complete response
    )
    
    final_response = response['message']['content']
    parsed_list = extract_list(final_response)
    
    # Return the complete response directly
    return parsed_list 


topic = "Python programming"
response = generate_topic_response(topic)

print(response)