ollama TUI chatbot

app.py

import os
import ollama
from rich.console import Console
from rich.markdown import Markdown
from rich.live import Live
from rich.panel import Panel
from rich.prompt import Prompt

# Limpa o terminal ao iniciar
os.system('clear' if os.name == 'posix' else 'cls')

console = Console()

def chat_loop():
    messages = []
    
    console.print(Panel.fit(
        "Chat Ollama + Rich Iniciado\n[grey]Pressione Ctrl+C para interromper uma resposta.\nDigite 'sair' para encerrar.",
        title="[bold magenta]IA Terminal[/]",
        border_style="magenta"
    ))

    while True:
        try:
            user_input = Prompt.ask("\n[bold cyan]Você[/]")
            
            if user_input.lower() in ["sair", "exit", "quit"]:
                break

            messages.append({'role': 'user', 'content': user_input})
            console.print("\n[bold green]IA:[/]")
            
            full_response = ""
            
            try:
                response = ollama.chat(
                    model='llama3.2:3b',
                    messages=messages,
                    stream=True,
                )

                with Live(Markdown(""), console=console, refresh_per_second=12, vertical_overflow="visible") as live:
                    for chunk in response:
                        content = chunk['message']['content']
                        full_response += content
                        live.update(Markdown(full_response))
                
                # Salva no histórico apenas se a resposta terminar normalmente
                messages.append({'role': 'assistant', 'content': full_response})

            except KeyboardInterrupt:
                # Interrompe apenas o streaming
                console.print("\n[yellow] Resposta interrompida pelo usuário.[/]")
                # Opcional: salvar o que foi gerado até agora no histórico
                if full_response:
                    messages.append({'role': 'assistant', 'content': f"{full_response} [INTERROMPIDO]"})

        except KeyboardInterrupt:
            # Caso o usuário dê Ctrl+C no prompt de entrada
            console.print("\n[bold red]Encerrando chat...[/]")
            break
        except Exception as e:
            console.print(f"[bold red]Erro:[/] {e}")
            break

if __name__ == "__main__":
    chat_loop()