Instruções de Sistema e Ferramentas para Agentes em Python

mkdir -p ~/projects/diagnostic-agent
cd ~/projects/diagnostic-agent
python3 -m venv venv
source venv/bin/activate
pip install --upgrade google-cloud-aiplatform[agent_engines,langchain] rich

import os
import vertexai
from vertexai import agent_engines
import osquery
from rich.console import Console
from rich.markdown import Markdown

PROJECT_ID = os.environ.get("GCP_PROJECT")
LOCATION = os.environ.get("GCP_REGION", "us-central1")
STAGING_BUCKET = os.environ.get("STAGING_BUCKET_URI")

vertexai.init(
    project=PROJECT_ID,
    location=LOCATION,
    staging_bucket=STAGING_BUCKET
)

MODEL = os.environ.get("GEMINI_MODEL", "gemini-2.0-flash")

instance = osquery.SpawnInstance()

def call_osquery(query: str):
    """Consulta o sistema operacional usando o osquery.
      
    Esta função envia uma query ao processo do osquery para retornar dados sobre a máquina atual, sistema operacional e processos em execução.
    Você também pode consultar o banco SQLite subjacente para inspecionar metadados usando tabelas de sistema como sqlite_master, sqlite_temp_master e tabelas virtuais.

    Args:
        query: str  Uma query SQL para uma das tabelas do osquery (ex: "select timestamp from time")

    Returns:
        ExtensionResponse: resposta do osquery com o status da requisição e o resultado da query.
    """
    if not instance.is_running():
        instance.open()

    result = instance.client.query(query)
    return result

def get_system_prompt():
    if not instance.is_running():
        instance.open()
    
    response = instance.client.query("select name from sqlite_temp_master").response
    tables = [ t["name"] for t in response ]
    return f"""
Papel:
  - Você é o agente de diagnóstico de emergência.
  - Você é o último recurso do usuário para diagnosticar problemas no computador.
  - Responda às consultas do usuário com o melhor de suas capacidades.
Ferramentas:
  - Você pode consultar o osquery usando a função call_osquery.
Contexto:
  - Utilize apenas tabelas desta lista: {tables}
  - Você pode descobrir schemas usando: PRAGMA table_info(tabela)
Tarefa:
  - Crie um plano de quais tabelas consultar para atender ao pedido do usuário.
  - Confirme o plano com o usuário antes de executar.
  - Se uma query falhar por coluna inexistente, execute a descoberta de schema e tente novamente.
  - Consulte a(s) tabela(s) necessária(s).
  - Apresente os resultados em formato legível para humanos (tabela ou lista).
    """

def main():
    agent = agent_engines.LangchainAgent(
        model = MODEL,
        system_instruction=get_system_prompt(),
        tools=[
            call_osquery,
        ],
    )
    
    console = Console()

    print("Bem-vindo ao Agente de Diagnóstico de Emergência\n")
    print("Qual é a natureza da sua emergência de diagnóstico?")

    while True:
        try:
            query = input(">> ")
        except EOFError:
            query = "exit"

        if query in ("exit", "quit"):
            break

        if query.strip() == "":
            continue
            
        response = agent.query(input=query)
        rendered_markdown = Markdown(response["output"])
        console.print(rendered_markdown)

    print("Até logo!")

if __name__ == "__main__":
    main()

$ python main.py
Bem-vindo ao Agente de Diagnóstico de Emergência

Qual é a natureza da sua emergência de diagnóstico?
>> 

Papel:
  - Você é o agente de diagnóstico de emergência.
  - Você é o último recurso do usuário para diagnosticar problemas no computador.
  - Responda às consultas do usuário com o melhor de suas capacidades.

Ferramentas:
  - Você pode consultar o osquery usando a função call_osquery.

Contexto:
  - Utilize apenas tabelas desta lista: {tables}
  - Você pode descobrir schemas usando: PRAGMA table_info(tabela)

Tarefa:
  - Crie um plano de quais tabelas consultar para atender ao pedido do usuário.
  - Confirme o plano com o usuário antes de executar.
  - Se uma query falhar por coluna inexistente, execute a descoberta de schema e tente novamente.
  - Consulte a(s) tabela(s) necessária(s).
  - Apresente os resultados em formato legível para humanos (tabela ou lista).

debug = False

def set_debug_mode(debug_mode: bool):
    """Ativa ou desativa o modo de depuração.
    
    Args:
        debug_mode (bool): True para ativar o debug, False para desativar.

    Returns:
        None
    """
    global debug
    debug = debug_mode

def call_osquery(query: str):
    """Consulta o sistema operacional usando o osquery."""
    if not instance.is_running():
        instance.open()

    if debug:
        print("Executando query: ", query)

    result = instance.client.query(query)
    if debug:
        print("Resultado da query: ", {
            "status": result.status.message if result.status else None, 
            "response": result.response if result.response else None
        })

    return result

$ python main.py
Bem-vindo ao Agente de Diagnóstico de Emergência

Qual é a natureza da sua emergência de diagnóstico?
>> execute um diagnóstico de nível 1 em modo debug
Executando query:  SELECT * FROM system_info
Resultado da query:  {'status': 'OK', 'response': [{...}]}
Executando query:  SELECT pid, name, user, cpu_percent FROM processes ORDER BY cpu_percent DESC LIMIT 10
Resultado da query:  {'status': 'no such column: user', 'response': None}
Executando query:  PRAGMA table_info(processes)
Resultado da query:  {'status': 'OK', 'response': [{'cid': '0', 'dflt_value': '', 'name': 'pid', 'notnull': '1', 'pk': '1', 'type': 'BIGINT'}, ...]}

from langchain_core.chat_history import InMemoryChatMessageHistory
import uuid

chats_by_session_id = {}

def get_chat_history(session_id: str) -> InMemoryChatMessageHistory:
    chat_history = chats_by_session_id.get(session_id)
    if chat_history is None:
        chat_history = InMemoryChatMessageHistory()
        chats_by_session_id[session_id] = chat_history
    return chat_history

response = agent.query(
    input=query,
    config={"configurable": {"session_id": session_id}}
)