import os
from dotenv import load_dotenv
from neo4j import GraphDatabase

# .env 파일로부터 접속 정보 로드
load_dotenv()

URI = os.getenv("NEO4J_URI")
USERNAME = os.getenv("NEO4J_USERNAME")
PASSWORD = os.getenv("NEO4J_PASSWORD")

# Neo4j 드라이버 인스턴스 생성
driver = GraphDatabase.driver(URI, auth=(USERNAME, PASSWORD))

try:
    # 데이터베이스 연결 통신 검증
    driver.verify_connectivity()
    print("Neo4j 연결 성공!")
finally:
    # 드라이버 자원 해제
    driver.close()

import os
import json
from typing import Literal

from dotenv import load_dotenv
from pydantic import BaseModel, Field

from langchain_groq import ChatGroq
from langchain_neo4j import Neo4jGraph


load_dotenv()

NEO4J_URI = os.getenv("NEO4J_URI")
NEO4J_USERNAME = os.getenv("NEO4J_USERNAME")
NEO4J_PASSWORD = os.getenv("NEO4J_PASSWORD")
GROQ_MODEL = os.getenv("GROQ_MODEL", "llama-3.3-70b-versatile")


# 1. Pydantic을 이용한 추출 스키마 정의
class KGNode(BaseModel):
    id: str = Field(description="노드 이름. 예: 김민수, 결제 시스템 리팩터링")
    type: Literal["Person", "Project", "Team", "Metric", "System", "Unknown"]


class KGRelationship(BaseModel):
    source: str = Field(description="출발 노드 id")
    target: str = Field(description="도착 노드 id")
    kind: Literal[
        "RESPONSIBLE_FOR",
        "AIMS_TO_REDUCE",
        "COLLABORATED_ON",
        "RELATED_TO",
    ]


class KGGraph(BaseModel):
    nodes: list[KGNode]
    relationships: list[KGRelationship]


# 2. Neo4jGraph 및 LLM 구조체 바인딩
graph = Neo4jGraph(
    url=NEO4J_URI,
    username=NEO4J_USERNAME,
    password=NEO4J_PASSWORD,
    database=os.getenv("NEO4J_DATABASE"),
)

llm = ChatGroq(
    model=GROQ_MODEL,
    temperature=0,
)

structured_llm = llm.with_structured_output(KGGraph)

# 3. 비정형 자연어 텍스트
text = """
김민수는 결제 시스템 리팩터링을 담당했다.
결제 시스템 리팩터링은 장애율을 낮추기 위한 프로젝트였다.
결제 시스템 리팩터링은 보안팀과 플랫폼팀이 공동으로 진행했다.
"""

prompt = f"""
다음 한국어 문장에서 지식 그래프를 추출하세요.

규칙:
- 문장에 명시된 정보만 사용하세요.
- 노드는 중요한 사람, 프로젝트, 팀, 지표, 시스템을 추출하세요.
- 관계는 가능한 한 아래 중 하나를 사용하세요:
  - RESPONSIBLE_FOR: 사람이 프로젝트나 업무를 담당함
  - AIMS_TO_REDUCE: 프로젝트가 어떤 지표를 낮추는 목적을 가짐
  - COLLABORATED_ON: 팀이 프로젝트를 공동 진행함
  - RELATED_TO: 위 관계로 표현하기 어려운 일반 관계
- relationship의 source와 target은 반드시 nodes의 id 중 하나여야 합니다.

문장:
{text}
"""

# LLM으로 구조화된 KG 추출 실행
kg = structured_llm.invoke(prompt)

print("LLM 추출 결과:")
print(json.dumps(kg.model_dump(), ensure_ascii=False, indent=2))


# 4. Neo4j DB 저장 (유니크 제약 조건 생성 및 UNWIND + MERGE 대량 저장)
graph.query("""
CREATE CONSTRAINT entity_id_unique IF NOT EXISTS
FOR (e:Entity)
REQUIRE e.id IS UNIQUE
""")

nodes = [node.model_dump() for node in kg.nodes]
relationships = [rel.model_dump() for rel in kg.relationships]

# 노드 멱등적 저장
graph.query(
    """
    UNWIND $nodes AS node
    MERGE (e:Entity {id: node.id})
    SET e.type = node.type
    """,
    params={"nodes": nodes},
)

# GraphCypherQAChain 조회를 돕기 위한 서브 라벨 동적 동기화
for node_type in ["Person", "Project", "Team", "Metric", "System", "Unknown"]:
    graph.query(
        f"""
        MATCH (e:Entity )
        SET e:{node_type}
        """,
        params={"node_type": node_type},
    )

# 관계 멱등적 저장
for kind in [
    "RESPONSIBLE_FOR",
    "AIMS_TO_REDUCE",
    "COLLABORATED_ON",
    "RELATED_TO",
]:
    graph.query(
        f"""
        UNWIND $relationships AS rel
        WITH rel
        WHERE rel.kind = $kind
        MATCH (source:Entity )
        MATCH (target:Entity )
        MERGE (source)-[r:{kind}]->(target)
        """,
        params={"relationships": relationships, "kind": kind},
    )

print("그래프 생성 완료!")

import os
import json
from typing import Literal

from dotenv import load_dotenv
from pydantic import BaseModel, Field
from pydantic_ai import Agent                # 👈 최신 Pydantic AI
from neo4j import GraphDatabase              # 👈 공식 Neo4j 드라이버

load_dotenv()

# 1. Pydantic 스키마 정의
class KGNode(BaseModel):
    id: str = Field(description="노드 이름")
    type: Literal["Person", "Project", "Team", "Metric", "System", "Unknown"]

class KGRelationship(BaseModel):
    source: str = Field(description="출발 노드 id")
    target: str = Field(description="도착 노드 id")
    kind: Literal["RESPONSIBLE_FOR", "AIMS_TO_REDUCE", "COLLABORATED_ON", "RELATED_TO"]

class KGGraph(BaseModel):
    nodes: list[KGNode]
    relationships: list[KGRelationship]

# 2. Graph Pruning 후처리 함수
def validate_kg(kg: KGGraph) -> KGGraph:
    node_ids = {node.id for node in kg.nodes}
    valid_relationships = [r for r in kg.relationships if r.source in node_ids and r.target in node_ids]
    return KGGraph(nodes=kg.nodes, relationships=valid_relationships)

# 3. Pydantic AI Agent 선언 (result_type에 KGGraph 지정)
kg_agent = Agent(
    'groq:llama-3.3-70b-versatile',
    result_type=KGGraph,  # 👈 반환받을 Pydantic 구조체 지정!
    system_prompt="""
    당신은 한국어 문장에서 노드와 관계를 파싱하는 지식 그래프 추출 AI입니다.
    규칙:
    - 노드는 사람, 프로젝트, 팀, 지표, 시스템을 추출하세요.
    - 관계 방향 규칙:
      - RESPONSIBLE_FOR: Person -> Project
      - AIMS_TO_REDUCE: Project -> Metric
      - COLLABORATED_ON: Team -> Project
    """
)

# 4. 실행 (단 한 줄로 타입 안정성이 보장된 결과 수신!)
text = """
김민수는 결제 시스템 리팩터링을 담당했다.
결제 시스템 리팩터링은 장애율을 낮추기 위한 프로젝트였다.
결제 시스템 리팩터링은 보안팀과 플랫폼팀이 공동으로 진행했다.
"""

result = kg_agent.run_sync(f"다음 문장에서 지식 그래프를 추출하세요:\n{text}")
clean_kg: KGGraph = validate_kg(result.data)  # 👈 result.data는 이미 KGGraph 객체!

print("Pydantic AI 추출 결과:")
print(json.dumps(clean_kg.model_dump(), ensure_ascii=False, indent=2))

# 5. Neo4j DB 저장 (공식 드라이버 사용)
driver = GraphDatabase.driver(os.getenv("NEO4J_URI"), auth=(os.getenv("NEO4J_USERNAME"), os.getenv("NEO4J_PASSWORD")))
nodes = [node.model_dump() for node in clean_kg.nodes]
relationships = [rel.model_dump() for rel in clean_kg.relationships]

with driver.session(database=os.getenv("NEO4J_DATABASE", "neo4j")) as session:
    session.run("CREATE CONSTRAINT entity_id_unique IF NOT EXISTS FOR (e:Entity) REQUIRE e.id IS UNIQUE")
    session.run(
        """
        UNWIND $nodes AS node
        MERGE (e:Entity {id: node.id})
        SET e.name = node.id, e.type = node.type
        """,
        nodes=nodes
    )
    for node_type in ["Person", "Project", "Team", "Metric", "System", "Unknown"]:
        session.run(f"MATCH (e:Entity ) SET e:{node_type}", node_type=node_type)
    for kind in ["RESPONSIBLE_FOR", "AIMS_TO_REDUCE", "COLLABORATED_ON", "RELATED_TO"]:
        session.run(
            """
            UNWIND $relationships AS rel
            WITH rel WHERE rel.kind = $kind
            MATCH (source:Entity {id: rel.source})
            MATCH (target:Entity {id: rel.target})
            MERGE (source)-[r:RESPONSIBLE_FOR|AIMS_TO_REDUCE|COLLABORATED_ON|RELATED_TO]->(target)
            """,
            relationships=relationships
        )

driver.close()
print("Pydantic AI 그래프 DB 구축 완료!")

import os
from dotenv import load_dotenv

from langchain_groq import ChatGroq
from langchain_neo4j import Neo4jGraph, GraphCypherQAChain

load_dotenv()

GROQ_MODEL = os.getenv("GROQ_MODEL", "llama-3.3-70b-versatile")

# Neo4jGraph 연결 설정
graph = Neo4jGraph(
    url=os.getenv("NEO4J_URI"),
    username=os.getenv("NEO4J_USERNAME"),
    password=os.getenv("NEO4J_PASSWORD"),
    database=os.getenv("NEO4J_DATABASE"),
)

# DB 스키마 최신화
graph.refresh_schema()

llm = ChatGroq(
    model=GROQ_MODEL,
    temperature=0,
)

# GraphCypherQAChain 생성
chain = GraphCypherQAChain.from_llm(
    llm=llm,
    graph=graph,
    verbose=True,           # 디버깅용 내부 체인 실행 로깅
    validate_cypher=True,   # LLM 생성 Cypher 쿼리 문법 검증
    allow_dangerous_requests=True,
)

print("Graph schema:")
print(graph.schema)

questions = [
    "김민수와 보안팀은 어떤 관계야?",
]

for question in questions:
    print("\n질문:", question)
    result = chain.invoke({"query": question})
    print("답변:", result["result"])

import os
from dotenv import load_dotenv
from neo4j import GraphDatabase
from pydantic_ai import Agent, RunContext

load_dotenv()

# Neo4j 파이썬 드라이버 수동 연결
driver = GraphDatabase.driver(
    os.getenv("NEO4J_URI"),
    auth=(os.getenv("NEO4J_USERNAME"), os.getenv("NEO4J_PASSWORD"))
)

# Pydantic AI Agent 생성
graph_agent = Agent(
    'groq:llama-3.3-70b-versatile',
    system_prompt="""
    당신은 Neo4j 그래프 데이터베이스 기반 질문 응답 AI 에이전트입니다.
    1. 먼저 get_db_schema 툴을 호출하여 DB 노드 라벨과 관계 스키마를 확인하세요.
    2. 질문에 필요한 데이터를 가져오기 위해 execute_cypher 툴로 READ(MATCH) 쿼리를 실행하세요.
    3. 조회된 결과를 바탕으로 사용자에게 자연스러운 답변을 제공하세요.
    """
)

@graph_agent.tool
def get_db_schema(ctx: RunContext) -> str:
    """Neo4j 그래프 데이터베이스의 스키마 메타데이터를 반환합니다."""
    return "Node Labels: Person, Project, Team, Metric\nRelationships: RESPONSIBLE_FOR, AIMS_TO_REDUCE, COLLABORATED_ON"

@graph_agent.tool
def execute_cypher(ctx: RunContext, cypher_query: str) -> list[dict]:
    """Neo4j 데이터베이스에 조회를 위한 MATCH Cypher 쿼리를 실행합니다."""
    # 파괴적/쓰기 쿼리 방지 안전 검증
    forbidden_keywords = ["DELETE", "DETACH", "DROP", "CREATE", "SET"]
    if any(kw in cypher_query.upper() for kw in forbidden_keywords):
        raise ValueError("안전을 위해 쓰기 및 삭제 쿼리는 허용되지 않습니다.")
    
    with driver.session() as session:
        result = session.run(cypher_query)
        return [record.data() for record in result]

# 에지 실행
response = graph_agent.run_sync("김민수와 보안팀은 어떤 관계야?")
print("최종 Pydantic AI 답변:", response.data)