""" Gemini API 연동 모듈 (좌표 추출 기능 추가) Google Gemini API를 사용하여 이미지와 텍스트 좌표를 함께 분석합니다. """ import base64 import logging import json from google import genai from google.genai import types from typing import Optional, Dict, Any, List from config import Config # 로깅 설정 logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) # --- 새로운 스키마 정의 --- # 좌표를 포함하는 값을 위한 재사용 가능한 스키마 ValueWithCoords = types.Schema( type=types.Type.OBJECT, properties={ "value": types.Schema(type=types.Type.STRING, description="추출된 텍스트 값"), "x": types.Schema(type=types.Type.NUMBER, description="텍스트의 시작 x 좌표"), "y": types.Schema(type=types.Type.NUMBER, description="텍스트의 시작 y 좌표"), }, required=["value", "x", "y"] ) # 범용 표의 한 행을 위한 스키마 GenericTableRow = types.Schema( type=types.Type.ARRAY, items=ValueWithCoords, description="표의 한 행. 각 셀의 값과 좌표를 포함합니다." ) # 페이지에서 추출된 범용 표를 위한 스키마 GenericTable = types.Schema( type=types.Type.OBJECT, properties={ "table_title": types.Schema(type=types.Type.STRING, description="추출된 표의 내용을 설명하는 제목 (예: '범례', 'IP 정보', '개정 이력')."), "table_data": types.Schema( type=types.Type.ARRAY, items=GenericTableRow, description="표의 데이터. 각 내부 리스트가 하나의 행을 나타냅니다." ) }, description="도면에서 발견된 구조화된 정보 블록이나 표." ) # 모든 필드가 ValueWithCoords를 사용하도록 스키마 업데이트 SCHEMA_EXPRESSWAY = types.Schema( type=types.Type.OBJECT, properties={ "도면명_line0": ValueWithCoords, "도면명_line1": ValueWithCoords, "도면명_line2": ValueWithCoords, "편철번호": ValueWithCoords, "도면번호": ValueWithCoords, "Main_Title": ValueWithCoords, "Sub_Title": ValueWithCoords, "수평_도면_축척": ValueWithCoords, "수직_도면_축척": ValueWithCoords, "적용표준버전": ValueWithCoords, "사업명_top": ValueWithCoords, "사업명_bot": ValueWithCoords, "시설_공구": ValueWithCoords, "설계공구_공구명": ValueWithCoords, "설계공구_범위": ValueWithCoords, "시공공구_공구명": ValueWithCoords, "시공공구_범위": ValueWithCoords, "건설분야": ValueWithCoords, "건설단계": ValueWithCoords, "설계사": ValueWithCoords, "시공사": ValueWithCoords, "노선이정": ValueWithCoords, "개정번호_1": ValueWithCoords, "개정날짜_1": ValueWithCoords, "개정내용_1": ValueWithCoords, "작성자_1": ValueWithCoords, "검토자_1": ValueWithCoords, "확인자_1": ValueWithCoords, "additional_tables": types.Schema( type=types.Type.ARRAY, items=GenericTable, description="도면에서 발견된 추가적인 표나 정보 블록 목록." ) }, ) SCHEMA_TRANSPORTATION = types.Schema( type=types.Type.OBJECT, properties={ "도면명": ValueWithCoords, "편철번호": ValueWithCoords, "도면번호": ValueWithCoords, "Main Title": ValueWithCoords, "Sub Title": ValueWithCoords, "수평축척": ValueWithCoords, "수직축척": ValueWithCoords, "적용표준": ValueWithCoords, "사업명": ValueWithCoords, "시설_공구": ValueWithCoords, "건설분야": ValueWithCoords, "건설단계": ValueWithCoords, "개정차수": ValueWithCoords, "개정일자": ValueWithCoords, "과업책임자": ValueWithCoords, "분야별책임자": ValueWithCoords, "설계자": ValueWithCoords, "위치정보": ValueWithCoords, "additional_tables": types.Schema( type=types.Type.ARRAY, items=GenericTable, description="도면에서 발견된 추가적인 표나 정보 블록 목록." ) }, ) class GeminiAnalyzer: """Gemini API 이미지 및 텍스트 분석 클래스""" def __init__(self, api_key: Optional[str] = None, model: Optional[str] = None): self.api_key = api_key or Config.GEMINI_API_KEY self.model = model or Config.GEMINI_MODEL self.default_prompt = Config.DEFAULT_PROMPT if not self.api_key: raise ValueError("Gemini API 키가 설정되지 않았습니다.") try: self.client = genai.Client(api_key=self.api_key) logger.info(f"Gemini 클라이언트 초기화 완료 (모델: {self.model})") except Exception as e: logger.error(f"Gemini 클라이언트 초기화 실패: {e}") raise def _get_schema(self, organization_type: str) -> types.Schema: """조직 유형에 따른 스키마를 반환합니다.""" return SCHEMA_EXPRESSWAY if organization_type == "한국도로공사" else SCHEMA_TRANSPORTATION def analyze_pdf_page( self, base64_data: str, text_blocks: List[Dict[str, Any]], prompt: Optional[str] = None, mime_type: str = "image/png", organization_type: str = "transportation" ) -> Optional[str]: """ Base64 이미지와 추출된 텍스트 좌표를 함께 분석합니다. Args: base64_data: Base64로 인코딩된 이미지 데이터. text_blocks: PDF에서 추출된 텍스트와 좌표 정보 리스트. prompt: 분석 요청 텍스트 (None인 경우 기본값 사용). mime_type: 이미지 MIME 타입. organization_type: 조직 유형 ("transportation" 또는 "expressway"). Returns: 분석 결과 JSON 문자열 또는 None (실패 시). """ try: # 텍스트 블록 정보를 JSON 문자열로 변환하여 프롬프트에 추가 text_context = "\n".join([ f"- text: '{block['text']}', bbox: ({block['bbox'][0]:.0f}, {block['bbox'][1]:.0f})" for block in text_blocks ]) analysis_prompt = ( (prompt or self.default_prompt) + "\n\n--- 추출된 텍스트와 좌표 정보 ---\n" + text_context + "\n\n--- 지시사항 ---\n" "1. 위 텍스트와 좌표 정보를 바탕으로, 이미지의 내용을 분석하여 JSON 스키마의 기본 필드를 채워주세요.\n" "2. **(중요)** 도면 내에 표나 사각형으로 구분된 정보 블록이 있다면, `additional_tables` 필드에 추가해주세요. 예를 들어 'IP 정보 표', '범례(Legend)', '구조물 설명', '개정 이력' 등이 해당됩니다.\n" " - 각 표/블록에 대해 `table_title`에 적절한 제목을 붙여주세요.\n" " - `table_data`에는 표의 모든 행을 추출하여 리스트 형태로 넣어주세요. 각 행은 셀들의 리스트입니다.\n" "3. 각 필드에 해당하는 텍스트를 찾고, 해당 텍스트의 'value'와 시작 'x', 'y' 좌표를 JSON에 기입하세요.\n" "4. 해당하는 값이 없으면 빈 문자열이나 빈 리스트를 사용하세요.\n" "\n--- 필드 설명 ---\n" "- `{ }_Title`: 중앙 상단의 비교적 큰 폰트입니다.\n" "- `사업명_top`에 해당하는 텍스트 아랫줄은 '시설_공구' 항목입니다.\n" "- `도면명_line{n}`: 도면명에 해당하는 여러 줄의 텍스트를 위에서부터 0, 1, 2 순서로 기입합니다. 만약 두 줄이라면 line0은 비워두고 line1, line2를 채웁니다.\n" "- `설계공구`/`시공공구`: '공구명'과 '범위'로 나뉘어 기입될 수 있습니다. (예: '설계공구 | 제2-1공구 | 12780.00-15860.00' -> `설계공구_공구명`: '제2-1공구', `설계공구_범위`: '12780.00-15860.00')\n" ) contents = [ types.Content( role="user", parts=[ types.Part.from_bytes( mime_type=mime_type, data=base64.b64decode(base64_data), ), types.Part.from_text(text=analysis_prompt), ], ) ] selected_schema = self._get_schema(organization_type) generate_content_config = types.GenerateContentConfig( temperature=0, top_p=0.05, response_mime_type="application/json", response_schema=selected_schema ) logger.info("Gemini API 분석 요청 시작 (텍스트 좌표 포함)...") response = self.client.models.generate_content( model=self.model, contents=contents, config=generate_content_config, ) if response and hasattr(response, 'text'): result = response.text # JSON 응답을 파싱하여 다시 직렬화 (일관된 포맷팅) parsed_json = json.loads(result) # 디버깅: Gemini 응답 내용 로깅 logger.info(f"=== Gemini 응답 디버깅 ===") logger.info(f"조직 유형: {organization_type}") logger.info(f"응답 필드 수: {len(parsed_json) if isinstance(parsed_json, dict) else 'N/A'}") if isinstance(parsed_json, dict): # 새로운 필드들이 응답에 포함되었는지 확인 new_fields = ["설계공구_Station_col1", "설계공구_Station_col2", "시공공구_Station_col1", "시공공구_Station_col2"] old_fields = ["설계공구_Station", "시공공구_Station"] logger.info("=== 새 필드 확인 ===") for field in new_fields: if field in parsed_json: field_data = parsed_json[field] if isinstance(field_data, dict) and field_data.get('value'): logger.info(f"✅ {field}: '{field_data.get('value')}' at ({field_data.get('x', 'N/A')}, {field_data.get('y', 'N/A')})") else: logger.info(f"⚠️ {field}: 빈 값 또는 잘못된 형식 - {field_data}") else: logger.info(f"❌ {field}: 응답에 없음") logger.info("=== 기존 필드 확인 ===") for field in old_fields: if field in parsed_json: field_data = parsed_json[field] if isinstance(field_data, dict) and field_data.get('value'): logger.info(f"⚠️ {field}: '{field_data.get('value')}' (기존 필드가 여전히 존재)") else: logger.info(f"⚠️ {field}: 빈 값 - {field_data}") else: logger.info(f"✅ {field}: 응답에 없음 (예상됨)") logger.info("=== 전체 응답 필드 목록 ===") for key in parsed_json.keys(): value = parsed_json[key] if isinstance(value, dict) and 'value' in value: logger.info(f"필드: {key} = '{value.get('value', '')}' at ({value.get('x', 'N/A')}, {value.get('y', 'N/A')})") else: logger.info(f"필드: {key} = {type(value).__name__}") logger.info("=== 디버깅 끝 ===") pretty_result = json.dumps(parsed_json, ensure_ascii=False, indent=2) logger.info(f"분석 완료: {len(pretty_result)} 문자") return pretty_result else: logger.error("API 응답에서 텍스트를 찾을 수 없습니다.") return None except Exception as e: logger.error(f"이미지 및 텍스트 분석 중 오류 발생: {e}") return None # --- 기존 다른 메서드들은 필요에 따라 수정 또는 유지 --- # analyze_image_stream_from_base64, analyze_pdf_images 등은 # 새로운 analyze_pdf_page 메서드와 호환되도록 수정 필요. # 지금은 핵심 기능에 집중. def validate_api_connection(self) -> bool: """API 연결 상태를 확인합니다.""" try: test_response = self.client.models.generate_content("안녕하세요") if test_response and hasattr(test_response, 'text'): logger.info("Gemini API 연결 테스트 성공") return True else: logger.error("Gemini API 연결 테스트 실패") return False except Exception as e: logger.error(f"Gemini API 연결 테스트 중 오류: {e}") return False