import ollama

import psycopg
from psycopg.rows import dict_row
from typing import List, Dict

def fetch_all_rows(conn_params: dict) -> List[Dict]:
    # Establish a connection to the PostgreSQL database
    conn = psycopg.connect(**conn_params, row_factory=dict_row)
    cursor = conn.cursor()
    
    # Define your SQL query to select all rows from the table
    sql_query = "SELECT * FROM public.primary_design_outcomes;"
    
    # Execute the query
    cursor.execute(sql_query)
    
    # Fetch all rows from the result set
    rows = cursor.fetchall()
    
    # Close the cursor and connection
    cursor.close()
    conn.close()
    
    return rows

# Example usage
conn_params = {
    "dbname": "aact_db",
    "user": "root",
    "password": "root",
    "host": "localhost",
    "port": "5432"
}

outcome_description = '''
Measure: {measure}
Time Frame: {time_frame}
Description: {description}
'''


if __name__ == "__main__":
    #check for model

    #get information
    rows_dicts = fetch_all_rows(conn_params)

    for row in rows_dicts[:3]:
        text_data = outcome_description.format(**row)
        r = ollama.generate(model='youainti/llama3.1-extractor:2024-08-28.2', prompt=text_data)

        print(text_data)
        print(r["response"])