from entities_extraction_agent import EntityExtractor, EntityType
from entities_extraction_agent.core.config import LLMConfig, ExtractionConfig

# 1. 定义要提取的实体类型
entity_types = [
    EntityType(
        name="Person",
        description="人名，包括姓名、职位等",
        examples=["张三", "李四", "王经理"]  # 可选：用于 few-shot 学习
    ),
    EntityType(
        name="Organization",
        description="组织机构名称",
        examples=["阿里巴巴", "清华大学"]
    ),
    EntityType(
        name="Location",
        description="地理位置",
        examples=["北京", "上海"]
    ),
]

# 2. 配置 LLM
llm_config = LLMConfig(
    provider="openai",
    model_name="gpt-3.5-turbo",
    temperature=0.0,
    api_key="your-api-key"  # 或使用环境变量 OPENAI_API_KEY
)

# 3. 配置提取策略
extraction_config = ExtractionConfig(
    strategy="few-shot",  # 或 "zero-shot"
    max_retries=3
)

# 4. 创建提取器
extractor = EntityExtractor(
    entity_types=entity_types,
    llm_config=llm_config,
    extraction_config=extraction_config
)

# 5. 提取实体
text = "马云创立了阿里巴巴集团，总部位于杭州。"
result = extractor.extract(text)

# 6. 查看结果
for entity in result.entities:
    print(f"{entity.text} ({entity.type})")

from entities_extraction_agent import EntityExtractor, EntityType
from entities_extraction_agent.preprocessing import TextPreprocessor
from entities_extraction_agent.postprocessing import EntityPostprocessor

# 预处理
preprocessor = TextPreprocessor(
    remove_html=True,
    normalize_whitespace=True
)

# 提取
extractor = EntityExtractor(entity_types=entity_types)
cleaned_text = preprocessor.preprocess(raw_text)
result = extractor.extract(cleaned_text)

# 后处理
postprocessor = EntityPostprocessor(
    deduplicate=True,
    normalize=True
)
final_result = postprocessor.postprocess(result)

# 查看统计
counts = postprocessor.get_entity_counts(final_result)
print(counts)  # Counter({'Person': 3, 'Organization': 2, ...})

from entities_extraction_agent.evaluation import EntityEvaluator
from entities_extraction_agent import Entity, ExtractionResult

# 准备真实标注
ground_truth = ExtractionResult(
    text=text,
    entities=[
        Entity(text="马云", type="Person"),
        Entity(text="阿里巴巴集团", type="Organization"),
        Entity(text="杭州", type="Location"),
    ]
)

# 评估
evaluator = EntityEvaluator()
metrics = evaluator.evaluate(predicted_result, ground_truth)

print(f"Precision: {metrics.precision:.3f}")
print(f"Recall: {metrics.recall:.3f}")
print(f"F1 Score: {metrics.f1_score:.3f}")

from entities_extraction_agent import EntityExtractor, EntityType, RelationshipType

# 定义实体类型
entity_types = [
    EntityType(name="Person", description="人名"),
    EntityType(name="Organization", description="组织名称"),
    EntityType(name="Location", description="地理位置"),
]

# 定义关系类型
relationship_types = [
    RelationshipType(
        name="works_for",
        description="雇佣关系，人与工作的组织",
        entity_types={"source": "Person", "target": "Organization"}
    ),
    RelationshipType(
        name="located_in",
        description="位置关系，组织所在的地理位置",
        entity_types={"source": "Organization", "target": "Location"}
    ),
]

# 创建提取器（同时提取实体和关系）
extractor = EntityExtractor(
    entity_types=entity_types,
    relationship_types=relationship_types,
)

text = "张三在阿里巴巴工作，公司总部位于杭州。"
result = extractor.extract(text)

# 查看提取的实体
for entity in result.entities:
    print(f"实体: {entity.text} ({entity.type})")

# 查看提取的关系
for rel in result.relationships:
    print(f"关系: {rel.source} -[{rel.type}]-> {rel.target}")

from entities_extraction_agent import DocumentExtractionPipeline, EntityType, RelationshipType

# 定义要提取的实体和关系类型
entity_types = [
    EntityType(name="Person", description="人名"),
    EntityType(name="Organization", description="组织名称"),
]

relationship_types = [
    RelationshipType(name="works_for", description="雇佣关系"),
]

# 创建流水线
pipeline = DocumentExtractionPipeline(
    entity_types=entity_types,
    relationship_types=relationship_types,
)

# 从 URL 处理文档
result = pipeline.from_url("https://example.com/document.pdf")

# 或从本地文件处理
# result = pipeline.from_file("./local_document.pdf")

# 或从纯文本处理
# result = pipeline.from_text("文档文本内容...")

# 查看结果
print(f"文档长度: {len(result.document_text)} 字符")
print(f"提取实体数: {len(result.entities)}")
print(f"提取关系数: {len(result.relationships)}")

for entity in result.entities:
    print(f"  - [{entity.type}] {entity.text}")

for rel in result.relationships:
    print(f"  - [{rel.type}] {rel.source} -> {rel.target}")