🔍 투자 리서치 에이전트 RAG 패턴

import requests
import os
import zipfile
import io
from pathlib import Path
from dotenv import load_dotenv

load_dotenv()
DART_KEY = os.getenv("DART_API_KEY")

class DartCollector:
    BASE = "https://opendart.fss.or.kr/api"

    def get_corp_code(self, company_name: str) -> str | None:
        """기업명 → DART 고유번호"""
        res = requests.get(
            f"{self.BASE}/company.json",
            params={"crtfc_key": DART_KEY, "corp_name": company_name},
        )
        items = res.json().get("list", [])
        return items[0]["corp_code"] if items else None

    def get_financial_report(self, corp_code: str, year: int, report_type: str = "11011") -> list[dict]:
        """
        사업보고서 주요 재무 항목 조회
        report_type: 11011=사업보고서, 11012=반기보고서, 11013=분기보고서
        """
        res = requests.get(
            f"{self.BASE}/fnlttSinglAcnt.json",
            params={
                "crtfc_key": DART_KEY,
                "corp_code": corp_code,
                "bsns_year": str(year),
                "reprt_code": report_type,
            },
        )
        data = res.json()
        if data.get("status") != "000":
            return []
        return data.get("list", [])

    def get_disclosures(self, corp_code: str, bgn_date: str, end_date: str) -> list[dict]:
        """공시 목록 조회"""
        res = requests.get(
            f"{self.BASE}/list.json",
            params={
                "crtfc_key": DART_KEY,
                "corp_code": corp_code,
                "bgn_de": bgn_date,
                "end_de": end_date,
                "sort": "date",
                "sort_mth": "desc",
                "count": 20,
            },
        )
        return res.json().get("list", [])

    def get_disclosure_text(self, rcept_no: str) -> str:
        """공시 원문 다운로드 및 텍스트 추출"""
        res = requests.get(
            f"{self.BASE}/document.xml",
            params={"crtfc_key": DART_KEY, "rcept_no": rcept_no},
        )
        # ZIP 해제 후 텍스트 추출
        try:
            z = zipfile.ZipFile(io.BytesIO(res.content))
            texts = []
            for name in z.namelist():
                if name.endswith(".xml") or name.endswith(".html"):
                    content = z.read(name).decode("utf-8", errors="ignore")
                    # 태그 제거
                    import re
                    texts.append(re.sub(r"<[^>]+>", " ", content))
            return " ".join(texts)[:10000]  # 10,000자 제한
        except Exception:
            return ""

def build_dart_documents(company_name: str, years: list[int]) -> list[dict]:
    """
    기업의 사업보고서 데이터를 문서 리스트로 변환
    반환: [{"content": str, "metadata": dict}]
    """
    collector = DartCollector()
    corp_code = collector.get_corp_code(company_name)
    if not corp_code:
        return []

    documents = []
    for year in years:
        items = collector.get_financial_report(corp_code, year)
        for item in items:
            content = (
                f"[{company_name} {year}년 재무데이터]\n"
                f"계정명: {item.get('account_nm')}\n"
                f"당기: {item.get('thstrm_amount')}원\n"
                f"전기: {item.get('frmtrm_amount')}원\n"
            )
            documents.append({
                "content": content,
                "metadata": {
                    "company": company_name,
                    "year": year,
                    "account": item.get("account_nm"),
                    "source": "DART",
                },
            })
    return documents

import requests
from langchain_core.tools import tool

SEC_HEADERS = {"User-Agent": "aidevops contact@aidevops.kr"}

class SecCollector:
    BASE = "https://data.sec.gov"

    def get_cik(self, ticker: str) -> str | None:
        """티커 → CIK 번호"""
        res = requests.get(
            "https://www.sec.gov/files/company_tickers.json",
            headers=SEC_HEADERS,
        )
        data = res.json()
        ticker_upper = ticker.upper()
        for item in data.values():
            if item["ticker"] == ticker_upper:
                return str(item["cik_str"]).zfill(10)
        return None

    def get_recent_10k(self, cik: str, count: int = 3) -> list[dict]:
        """최근 10-K(연간보고서) 목록"""
        res = requests.get(
            f"{self.BASE}/submissions/CIK{cik}.json",
            headers=SEC_HEADERS,
        )
        data = res.json()
        filings = data.get("filings", {}).get("recent", {})

        results = []
        forms = filings.get("form", [])
        dates = filings.get("filingDate", [])
        accs  = filings.get("accessionNumber", [])

        for i, form in enumerate(forms):
            if form == "10-K" and len(results) < count:
                results.append({
                    "form": form,
                    "date": dates[i],
                    "accession": accs[i],
                    "url": f"https://www.sec.gov/Archives/edgar/full-index/{dates[i][:4]}/",
                })
        return results

    def get_filing_text(self, cik: str, accession: str) -> str:
        """10-K 텍스트 추출 (앞 8,000자)"""
        acc_clean = accession.replace("-", "")
        url = f"{self.BASE}/Archives/edgar/data/{int(cik)}/{acc_clean}/{accession}.txt"
        try:
            res = requests.get(url, headers=SEC_HEADERS, timeout=10)
            import re
            text = re.sub(r"<[^>]+>", " ", res.text)
            text = re.sub(r"\s+", " ", text)
            return text[:8000]
        except Exception:
            return ""

def build_sec_documents(ticker: str) -> list[dict]:
    """미국 기업 10-K 데이터를 문서 리스트로 변환"""
    collector = SecCollector()
    cik = collector.get_cik(ticker)
    if not cik:
        return []

    documents = []
    for filing in collector.get_recent_10k(cik, count=2):
        text = collector.get_filing_text(cik, filing["accession"])
        if text:
            documents.append({
                "content": f"[{ticker} 10-K {filing['date']}]\n{text}",
                "metadata": {
                    "ticker": ticker,
                    "form": "10-K",
                    "date": filing["date"],
                    "source": "SEC",
                },
            })
    return documents

import chromadb
from chromadb.config import Settings
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import Chroma
from pathlib import Path

# 한국어 특화 임베딩 모델 (로컬, 무료)
EMBED_MODEL = "jhgan/ko-sroberta-multitask"
DB_PATH = "D:/ai-agent-finance/chroma_db"

def get_embeddings():
    return HuggingFaceEmbeddings(
        model_name=EMBED_MODEL,
        model_kwargs={"device": "cpu"},
        encode_kwargs={"normalize_embeddings": True},
    )

def get_vectorstore(collection: str = "finance") -> Chroma:
    return Chroma(
        collection_name=collection,
        embedding_function=get_embeddings(),
        persist_directory=DB_PATH,
    )

def ingest_documents(documents: list[dict], collection: str = "finance") -> int:
    """
    문서를 청킹해 벡터 DB에 저장합니다.
    documents: [{"content": str, "metadata": dict}]
    반환: 저장된 청크 수
    """
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=512,
        chunk_overlap=64,
        separators=["\n\n", "\n", "。", ". ", " "],
    )

    chunks, metadatas = [], []
    for doc in documents:
        splits = splitter.split_text(doc["content"])
        chunks.extend(splits)
        metadatas.extend([doc["metadata"]] * len(splits))

    vs = get_vectorstore(collection)
    vs.add_texts(texts=chunks, metadatas=metadatas)
    vs.persist()
    print(f"✓ {len(chunks)}개 청크를 '{collection}' 컬렉션에 저장했습니다.")
    return len(chunks)

# 사용 예시
if __name__ == "__main__":
    from collectors.dart_collector import build_dart_documents
    from collectors.sec_collector import build_sec_documents

    # 한국 기업 인제스트
    docs = build_dart_documents("삼성전자", years=[2023, 2024])
    ingest_documents(docs, collection="kr_finance")

    # 미국 기업 인제스트
    docs = build_sec_documents("AAPL")
    ingest_documents(docs, collection="us_finance")

from langchain_core.tools import tool
from rag.vectorstore import get_vectorstore

def _search(query: str, collection: str, k: int = 4) -> str:
    vs = get_vectorstore(collection)
    docs = vs.similarity_search(query, k=k)
    if not docs:
        return "관련 문서를 찾을 수 없습니다."
    return "\n\n---\n\n".join(
        f"[출처: {d.metadata.get('source','?')} {d.metadata.get('date','')}]\n{d.page_content}"
        for d in docs
    )

@tool
def search_kr_disclosure(query: str) -> str:
    """
    한국 기업 DART 공시 데이터베이스에서 관련 내용을 검색합니다.
    query: 검색할 내용 (예: "삼성전자 2024 영업이익", "SK하이닉스 반도체 재고")
    """
    return _search(query, "kr_finance")

@tool
def search_us_filing(query: str) -> str:
    """
    미국 기업 SEC 10-K 데이터베이스에서 관련 내용을 검색합니다.
    query: 검색할 내용 (예: "Apple revenue 2024", "NVDA data center risk")
    """
    return _search(query, "us_finance")

from langchain_ollama import ChatOllama
from langchain_anthropic import ChatAnthropic
from langchain.agents import create_react_agent, AgentExecutor
from langchain_core.prompts import PromptTemplate
from rag.retriever import search_kr_disclosure, search_us_filing
from tools.us_stock import get_us_stock_price
from tools.kr_stock import get_kr_stock_price
import os

RESEARCH_TOOLS = [
    search_kr_disclosure,
    search_us_filing,
    get_kr_stock_price,
    get_us_stock_price,
]

SYSTEM = """당신은 금융 투자 리서치 전문가입니다.
공시 문서와 재무 데이터를 기반으로 정확한 분석을 제공합니다.

중요 원칙:
- 검색된 문서의 내용만 근거로 답변합니다
- 출처(DART/SEC)를 명시합니다
- 수치는 원문 그대로 인용합니다
- 투자 권유는 하지 않습니다"""

def build_research_agent(use_claude: bool = False) -> AgentExecutor:
    if use_claude:
        llm = ChatAnthropic(
            model="claude-sonnet-4-6",
            api_key=os.getenv("ANTHROPIC_API_KEY"),
            temperature=0,
        )
    else:
        llm = ChatOllama(model="llama3.1:8b", temperature=0)

    prompt = PromptTemplate.from_template(
        SYSTEM + """

사용 가능한 도구:
{tools}

도구 이름: {tool_names}

질문: {input}
{agent_scratchpad}"""
    )

    agent = create_react_agent(llm, RESEARCH_TOOLS, prompt)
    return AgentExecutor(
        agent=agent,
        tools=RESEARCH_TOOLS,
        verbose=True,
        max_iterations=6,
        handle_parsing_errors=True,
    )

from agents.research_agent import build_research_agent

agent = build_research_agent(use_claude=False)  # 로컬 Ollama

questions = [
    "삼성전자 2024년 영업이익은 얼마인가요?",
    "SK하이닉스의 주요 재무 리스크는 무엇인가요?",
    "Apple의 2024 10-K에서 언급된 주요 사업 리스크는?",
    "삼성전자 반도체 부문 실적 추이를 설명해주세요",
]

for q in questions:
    print(f"\nQ: {q}")
    result = agent.invoke({"input": q})
    print(f"A: {result['output']}")