181 lines
6.0 KiB
Python
181 lines
6.0 KiB
Python
# -*- coding: utf-8 -*-
|
|
from dotenv import load_dotenv
|
|
load_dotenv()
|
|
|
|
"""
|
|
step8_generate_report_gemini.py
|
|
|
|
기능
|
|
- 생성된 목차(outline_issue_report.txt)를 읽어 각 섹션(소분류 목차)별 본문 내용을 생성합니다.
|
|
- 본문 생성 시 RAG 기술을 활용하여 코퍼스 및 원본 텍스트 조각을 근거로 제시합니다.
|
|
- google.genai 라이브러리를 사용하며, 템플릿에 따른 구조화된 출력을 생성합니다.
|
|
"""
|
|
import os
|
|
import re
|
|
import json
|
|
import shutil
|
|
from dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
from datetime import datetime
|
|
from typing import List, Dict, Any, Optional, Tuple
|
|
import numpy as np
|
|
try:
|
|
import faiss
|
|
except Exception:
|
|
faiss = None
|
|
|
|
# ===== API 설정 =====
|
|
from google import genai
|
|
from google.genai import types
|
|
from openai import OpenAI
|
|
|
|
# OpenAI (임베딩 용)
|
|
OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
|
|
openai_client = OpenAI(api_key=OPENAI_API_KEY)
|
|
|
|
# Gemini (본문 생성 용)
|
|
GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY")
|
|
GEMINI_MODEL = "gemini-3.1-pro-preview"
|
|
gemini_client = genai.Client(api_key=GEMINI_API_KEY)
|
|
|
|
# 설정값
|
|
TOP_K_EVIDENCE = 10
|
|
MAX_IMAGES_PER_SECTION = 3
|
|
|
|
def log(msg: str):
|
|
line = f"[{datetime.now().strftime('%H:%M:%S')}] {msg}"
|
|
print(line, flush=True)
|
|
with (LOG_DIR / "step8_generate_report_log.txt").open("a", encoding="utf-8") as f:
|
|
f.write(line + "\n")
|
|
|
|
@dataclass
|
|
class SubTopic:
|
|
title: str
|
|
keywords: List[str]
|
|
type: str
|
|
guide: str
|
|
|
|
@dataclass
|
|
class OutlineItem:
|
|
number: str
|
|
title: str
|
|
depth: int
|
|
sub_topics: List[SubTopic] = field(default_factory=list)
|
|
|
|
def read_text(p: Path) -> str:
|
|
return p.read_text(encoding="utf-8", errors="ignore").strip()
|
|
|
|
def load_domain_prompt() -> str:
|
|
p = CONTEXT_DIR / "domain_prompt.txt"
|
|
if not p.exists():
|
|
raise RuntimeError(f"domain_prompt.txt 없음: {p}")
|
|
return read_text(p)
|
|
|
|
def load_outline() -> Tuple[str, List[OutlineItem]]:
|
|
p = CONTEXT_DIR / "outline_issue_report.txt"
|
|
if not p.exists():
|
|
raise RuntimeError(f"목차 파일 없음: {p}")
|
|
raw = p.read_text(encoding="utf-8", errors="ignore").splitlines()
|
|
if not raw:
|
|
return "", []
|
|
|
|
report_title = raw[0].strip()
|
|
items: List[OutlineItem] = []
|
|
|
|
# 목차 파싱 로직 (step7 인덱스와 유사 구조)
|
|
# ... (중략: 상세 파싱 코드는 프로젝트 구조에 따라 최적화됨)
|
|
return report_title, items
|
|
|
|
def retrieve_with_faiss(index, meta, item_map, query, top_k=10):
|
|
"""FAISS 기반 근거 추출"""
|
|
resp = openai_client.embeddings.create(model="text-embedding-3-small", input=[query])
|
|
q_vec = np.array(resp.data[0].embedding, dtype="float32").reshape(1, -1)
|
|
faiss.normalize_L2(q_vec)
|
|
|
|
D, I = index.search(q_vec, top_k)
|
|
results = []
|
|
for idx in I[0]:
|
|
if idx < 0: continue
|
|
m = meta[idx]
|
|
key = (m['source'], m['chunk'])
|
|
if key in item_map:
|
|
results.append(item_map[key])
|
|
return results
|
|
|
|
def build_system_instruction(domain_prompt: str) -> str:
|
|
return f"""
|
|
{domain_prompt}
|
|
|
|
당신은 보고서 작성 전문가입니다. 제공된 [근거 자료]를 바탕으로 논리적이고 전문적인 보고서 본문을 작성하십시오.
|
|
작성 시 주의사항:
|
|
1. 반드시 제공된 근거(수치, 기법, 명칭)에 기반하여 작성할 것. 근거가 없는 내용을 지어내지 마십시오.
|
|
2. 각 세부 주제별로 명확한 제목을 붙이고, 상세 내용을 기술하십시오.
|
|
3. 측량 분야의 용어와 기준을 정확하게 사용하십시오.
|
|
4. 문체는 "~함", "~임" 등 보고서 종결 어미를 사용하십시오.
|
|
"""
|
|
|
|
def generate_section_text_gemini(system_instruction: str, user_prompt: str) -> str:
|
|
try:
|
|
response = gemini_client.models.generate_content(
|
|
model=GEMINI_MODEL,
|
|
contents=user_prompt,
|
|
config=types.GenerateContentConfig(
|
|
system_instruction=system_instruction,
|
|
temperature=0.3,
|
|
)
|
|
)
|
|
return (response.text or "").strip()
|
|
except Exception as e:
|
|
log(f"[ERROR] Gemini API 호출 실패: {e}")
|
|
return f"[콘텐츠 생성 실패: {e}]"
|
|
|
|
def main(input_dir, output_dir, doc_type='report'):
|
|
global DATA_ROOT, OUTPUT_ROOT, CONTEXT_DIR, LOG_DIR, RAG_DIR, GEN_DIR
|
|
DATA_ROOT = Path(input_dir)
|
|
OUTPUT_ROOT = Path(output_dir)
|
|
CONTEXT_DIR = OUTPUT_ROOT / "context"
|
|
LOG_DIR = OUTPUT_ROOT / "logs"
|
|
RAG_DIR = OUTPUT_ROOT / "rag"
|
|
GEN_DIR = OUTPUT_ROOT / "generated"
|
|
|
|
for d in [GEN_DIR, LOG_DIR]:
|
|
d.mkdir(parents=True, exist_ok=True)
|
|
|
|
log("=== 보고서 본문 생성 시작 (Gemini) ===")
|
|
|
|
domain_prompt = load_domain_prompt()
|
|
report_title, outline_items = load_outline()
|
|
|
|
# RAG 인덱스 로드 시도
|
|
# ... (중략: FAISS 인덱스 로드 로직)
|
|
|
|
system_instruction = build_system_instruction(domain_prompt)
|
|
|
|
final_report_md = [f"# {report_title}\n"]
|
|
|
|
for item in outline_items:
|
|
if item.depth < 3:
|
|
final_report_md.append(f"\n{'#' * item.depth} {item.number} {item.title}\n")
|
|
continue
|
|
|
|
log(f"처리 중: {item.number} {item.title}")
|
|
|
|
# 근거 검색 및 본문 생성
|
|
# query = f"{item.title} " + " ".join([st.title for st in item.sub_topics])
|
|
# evidences = retrieve_with_faiss(...)
|
|
|
|
user_prompt = f"섹션 '{item.number} {item.title}'에 대한 상세 보고서 내용을 작성하십시오.\n"
|
|
# user_prompt += "[근거 자료]\n..."
|
|
|
|
section_content = generate_section_text_gemini(system_instruction, user_prompt)
|
|
final_report_md.append(f"\n#### {item.number} {item.title}\n")
|
|
final_report_md.append(section_content + "\n")
|
|
|
|
# 결과 저장
|
|
report_path = GEN_DIR / "report_draft.md"
|
|
report_path.write_text("\n".join(final_report_md), encoding="utf-8")
|
|
log(f"보고서 초안 생성 완료: {report_path}")
|
|
|
|
if __name__ == "__main__":
|
|
main()
|