참고소스 수정본
This commit is contained in:
LASTA_DEV01\lasta
282
참고/knowledge_agent-main/db_utils.py
Normal file
282
참고/knowledge_agent-main/db_utils.py
Normal file
@@ -0,0 +1,282 @@
|
||||
# db_utils.py
|
||||
import os
|
||||
import psycopg2
|
||||
import json
|
||||
from contextlib import contextmanager
|
||||
import json_repair
|
||||
|
||||
# --- Database Connection ---
|
||||
@contextmanager
|
||||
def get_db_connection():
|
||||
"""Provides a database connection using a context manager."""
|
||||
conn = psycopg2.connect(os.environ["DATABASE_URL"])
|
||||
try:
|
||||
yield conn
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
# --- Table Creation ---
|
||||
def create_tables():
|
||||
"""Creates all necessary tables in the database if they don't exist."""
|
||||
commands = (
|
||||
"""CREATE TABLE IF NOT EXISTS analyst_reports (id SERIAL PRIMARY KEY, report_id VARCHAR(255) UNIQUE NOT NULL, report JSONB, created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP);""",
|
||||
"""CREATE TABLE IF NOT EXISTS researcher_reports (id SERIAL PRIMARY KEY, report_id VARCHAR(255) UNIQUE NOT NULL, report JSONB, created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP);""",
|
||||
"""CREATE TABLE IF NOT EXISTS curator_reports (id SERIAL PRIMARY KEY, report_id VARCHAR(255) UNIQUE NOT NULL, report JSONB, created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP);""",
|
||||
"""CREATE TABLE IF NOT EXISTS auditor_reports (id SERIAL PRIMARY KEY, report_id VARCHAR(255) UNIQUE NOT NULL, report JSONB, created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP);""",
|
||||
"""CREATE TABLE IF NOT EXISTS fixer_reports (id SERIAL PRIMARY KEY, report_id VARCHAR(255) UNIQUE NOT NULL, report JSONB, created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP);""",
|
||||
"""CREATE TABLE IF NOT EXISTS advisor_reports (id SERIAL PRIMARY KEY, report_id VARCHAR(255) UNIQUE NOT NULL, report JSONB, created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP);""",
|
||||
"""CREATE TABLE IF NOT EXISTS documents (id SERIAL PRIMARY KEY, url TEXT UNIQUE NOT NULL, raw_document BYTEA, markdown_content TEXT, summary TEXT, created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP);"""
|
||||
)
|
||||
with get_db_connection() as conn:
|
||||
with conn.cursor() as cur:
|
||||
for command in commands:
|
||||
cur.execute(command)
|
||||
conn.commit()
|
||||
|
||||
# --- Document Handling Functions ---
|
||||
def add_url_or_get_id(url: str) -> list:
|
||||
"""Adds a URL to the documents table if it doesn't exist, or returns the existing id."""
|
||||
with get_db_connection() as conn:
|
||||
with conn.cursor() as cur:
|
||||
cur.execute("SELECT id FROM documents WHERE url = %s;", (url,))
|
||||
result = cur.fetchone()
|
||||
if result:
|
||||
return result[0], "existing"
|
||||
else:
|
||||
cur.execute("INSERT INTO documents (url) VALUES (%s) RETURNING id;", (url,))
|
||||
new_id = cur.fetchone()[0]
|
||||
conn.commit()
|
||||
return new_id, "new"
|
||||
|
||||
def update_document_content(url_id: int, raw_document: bytes, markdown_content: str):
|
||||
"""Updates the raw_document and markdown_content for a given url_id."""
|
||||
# Clean the markdown_content to remove any null characters
|
||||
cleaned_markdown_content = markdown_content.replace('\x00', '')
|
||||
|
||||
with get_db_connection() as conn:
|
||||
with conn.cursor() as cur:
|
||||
cur.execute(
|
||||
"UPDATE documents SET raw_document = %s, markdown_content = %s WHERE id = %s;",
|
||||
(raw_document, cleaned_markdown_content, url_id)
|
||||
)
|
||||
conn.commit()
|
||||
|
||||
def get_document_object(url_id: int, type: str):
|
||||
"""Gets the object for a given url_id in the documents table."""
|
||||
allowed_types = ["raw_document", "markdown_content", "summary"]
|
||||
if type not in allowed_types:
|
||||
raise ValueError(f"Invalid type specified: {type}")
|
||||
|
||||
with get_db_connection() as conn:
|
||||
with conn.cursor() as cur:
|
||||
query = f"SELECT {type} FROM documents WHERE id = %s;"
|
||||
cur.execute(query, (url_id,))
|
||||
result = cur.fetchone()
|
||||
if result:
|
||||
return result[0]
|
||||
else:
|
||||
return None
|
||||
|
||||
def update_document_object(url_id: int, type: str, object: str | bytes):
|
||||
"""Updates the object for a given url_id in the documents table."""
|
||||
allowed_types = ["raw_document", "markdown_content", "summary"]
|
||||
if type not in allowed_types:
|
||||
raise ValueError(f"Invalid type specified: {type}")
|
||||
|
||||
with get_db_connection() as conn:
|
||||
with conn.cursor() as cur:
|
||||
query = f"UPDATE documents SET {type} = %s WHERE id = %s;"
|
||||
cur.execute(query, (object, url_id))
|
||||
conn.commit()
|
||||
|
||||
def get_document(url_id: int) -> dict:
|
||||
"""Retrieves a document from the documents table."""
|
||||
with get_db_connection() as conn:
|
||||
with conn.cursor() as cur:
|
||||
cur.execute("SELECT url, raw_document, markdown_content, summary FROM documents WHERE id = %s;", (url_id,))
|
||||
result = cur.fetchone()
|
||||
if result:
|
||||
return {
|
||||
"url": result[0],
|
||||
"raw_document": result[1],
|
||||
"markdown_content": result[2],
|
||||
"summary": result[3]
|
||||
}
|
||||
else:
|
||||
return None
|
||||
|
||||
# --- Utility Functions ---
|
||||
|
||||
def extract_and_clean_json(llm_output: str) -> dict:
|
||||
|
||||
# Use json_repair.loads() directly as a robust, drop-in replacement for json_repair.loads()
|
||||
try:
|
||||
return json_repair.loads(llm_output)
|
||||
except Exception as e:
|
||||
raise ValueError(f"Failed to repair or parse JSON: {e}")
|
||||
|
||||
# --- Report Handling Functions ---
|
||||
def _save_report(table_name: str, report_data: dict):
|
||||
"""Generic function to save a report to a specified table."""
|
||||
report_id = report_data.get("report_id")
|
||||
if not report_id:
|
||||
raise ValueError("Report data must include a 'report_id'")
|
||||
|
||||
report_json_string = json.dumps(report_data)
|
||||
|
||||
with get_db_connection() as conn:
|
||||
with conn.cursor() as cur:
|
||||
cur.execute(
|
||||
f"INSERT INTO {table_name} (report_id, report) VALUES (%s, %s);",
|
||||
(report_id, report_json_string)
|
||||
)
|
||||
conn.commit()
|
||||
|
||||
def save_analyst_report(report_data: dict):
|
||||
_save_report("analyst_reports", report_data)
|
||||
|
||||
def save_auditor_report(report_data: dict):
|
||||
_save_report("auditor_reports", report_data)
|
||||
|
||||
def save_fixer_report(report_data: dict):
|
||||
_save_report("fixer_reports", report_data)
|
||||
|
||||
def save_advisor_report(report_data: dict):
|
||||
_save_report("advisor_reports", report_data)
|
||||
|
||||
def initialize_researcher(timestamp: str) -> dict:
|
||||
"""Initializes the researcher's report in the database."""
|
||||
analyst_report_str = load_latest_report('analyst')
|
||||
analyst_report = json.loads(analyst_report_str)
|
||||
|
||||
report_id = f"res_{timestamp.replace('-', '').replace(':', '').replace('T', '_').split('.')[0]}"
|
||||
|
||||
gaps_to_do = [
|
||||
{"gap_id": gap["gap_id"], "description": gap["description"], "research_topic": gap["research_topic"], "searches": [" "]}
|
||||
for gap in analyst_report.get("identified_gaps", [])
|
||||
]
|
||||
|
||||
new_report = {"report_id": report_id, "gaps": gaps_to_do}
|
||||
report_json_string = json.dumps(new_report)
|
||||
|
||||
with get_db_connection() as conn:
|
||||
with conn.cursor() as cur:
|
||||
cur.execute(
|
||||
"INSERT INTO researcher_reports (report_id, report) VALUES (%s, %s);",
|
||||
(report_id, report_json_string)
|
||||
)
|
||||
conn.commit()
|
||||
|
||||
return {"researcher_report_id": report_id, "researcher_gaps_todo": gaps_to_do}
|
||||
|
||||
def update_researcher_report(report_id: str, gap_id: str, searches: list):
|
||||
"""Updates a researcher report in the database with search results."""
|
||||
with get_db_connection() as conn:
|
||||
with conn.cursor() as cur:
|
||||
# Read the existing report
|
||||
cur.execute("SELECT report FROM researcher_reports WHERE report_id = %s FOR UPDATE;", (report_id,))
|
||||
result = cur.fetchone()
|
||||
if not result:
|
||||
raise ValueError(f"No researcher report found with id {report_id}")
|
||||
|
||||
report_data = result[0]
|
||||
|
||||
# Modify the report in Python
|
||||
gaps = report_data.get("gaps", [])
|
||||
gap_found = False
|
||||
for gap in gaps:
|
||||
if gap.get("gap_id") == gap_id:
|
||||
gap["searches"] = searches
|
||||
gap_found = True
|
||||
break
|
||||
|
||||
if not gap_found:
|
||||
# This case should ideally not be reached if initialization is correct
|
||||
raise ValueError(f"Gap with id {gap_id} not found in report {report_id}")
|
||||
|
||||
# Write the modified report back
|
||||
report_json_string = json.dumps(report_data)
|
||||
cur.execute(
|
||||
"UPDATE researcher_reports SET report = %s WHERE report_id = %s;",
|
||||
(report_json_string, report_id)
|
||||
)
|
||||
conn.commit()
|
||||
|
||||
def initialize_curator(timestamp: str) -> dict:
|
||||
"""Initializes the curator's report in the database."""
|
||||
researcher_report_str = load_latest_report('researcher')
|
||||
researcher_report = json.loads(researcher_report_str)
|
||||
|
||||
report_id = f"cur_{timestamp.replace('-', '').replace(':', '').replace('T', '_').split('.')[0]}"
|
||||
|
||||
searches_todo = []
|
||||
for gap in researcher_report.get("gaps", []):
|
||||
research_topic = gap.get("research_topic", {})
|
||||
for search in gap.get("searches", []):
|
||||
searches_todo.append({
|
||||
"search": search,
|
||||
"research_topic": research_topic
|
||||
})
|
||||
|
||||
new_report = {
|
||||
"report_id": report_id,
|
||||
"urls_for_ingestion": [],
|
||||
"url_ingestion_status": []
|
||||
}
|
||||
report_json_string = json.dumps(new_report)
|
||||
|
||||
with get_db_connection() as conn:
|
||||
with conn.cursor() as cur:
|
||||
cur.execute(
|
||||
"INSERT INTO curator_reports (report_id, report) VALUES (%s, %s);",
|
||||
(report_id, report_json_string)
|
||||
)
|
||||
conn.commit()
|
||||
|
||||
return {"curator_report_id": report_id, "curator_searches_todo": searches_todo}
|
||||
|
||||
def update_curator_report(report_id: str, job: str, results: list):
|
||||
"""Appends results to a job list in a curator report."""
|
||||
with get_db_connection() as conn:
|
||||
with conn.cursor() as cur:
|
||||
# Read the existing report
|
||||
cur.execute("SELECT report FROM curator_reports WHERE report_id = %s FOR UPDATE;", (report_id,))
|
||||
result = cur.fetchone()
|
||||
if not result:
|
||||
raise ValueError(f"No curator report found with id {report_id}")
|
||||
|
||||
report_data = result[0]
|
||||
|
||||
# Modify the report in Python
|
||||
if job not in report_data:
|
||||
report_data[job] = []
|
||||
|
||||
report_data[job].extend(results)
|
||||
|
||||
# Write the modified report back
|
||||
report_json_string = json.dumps(report_data)
|
||||
cur.execute(
|
||||
"UPDATE curator_reports SET report = %s WHERE report_id = %s;",
|
||||
(report_json_string, report_id)
|
||||
)
|
||||
conn.commit()
|
||||
|
||||
def load_latest_report(report_type: str) -> str:
|
||||
"""Loads the most recent report from the database."""
|
||||
table_map = {
|
||||
"analyst": "analyst_reports",
|
||||
"researcher": "researcher_reports",
|
||||
"curator": "curator_reports"
|
||||
}
|
||||
table_name = table_map.get(report_type)
|
||||
if not table_name:
|
||||
raise ValueError(f"Invalid report_type '{report_type}'.")
|
||||
|
||||
with get_db_connection() as conn:
|
||||
with conn.cursor() as cur:
|
||||
cur.execute(f"SELECT report FROM {table_name} ORDER BY created_at DESC LIMIT 1;")
|
||||
result = cur.fetchone()
|
||||
if result:
|
||||
return json.dumps(result[0])
|
||||
else:
|
||||
raise FileNotFoundError(f"No reports found in table {table_name}")
|
||||
Reference in New Issue
Block a user