agent.py

import os
from typing import List, Dict, Any
import gradio as gr
from dotenv import load_dotenv
from transformers import pipeline

# 環境変数の読み込み
load_dotenv()

class ContractAgent:
    def __init__(self):
        self.name = "契約管理エージェント"
        self.description = "法人契約の管理をサポートするエージェント"
        self.hf_token = os.getenv("HUGGINGFACE_TOKEN")
        
        # Hugging Faceのパイプラインの初期化
        self.text_classifier = pipeline(
            "text-classification",
            model="cl-tohoku/bert-base-japanese",
            token=self.hf_token
        )
        
        self.ner = pipeline(
            "token-classification",
            model="cl-tohoku/bert-base-japanese-ner",
            token=self.hf_token
        )
        
    def process_contract(self, contract_text: str) -> Dict[str, Any]:
        """
        契約書のテキストを処理する関数
        
        Args:
            contract_text (str): 契約書のテキスト
            
        Returns:
            Dict[str, Any]: 処理結果
        """
        # テキスト分類の実行
        classification = self.text_classifier(contract_text)
        
        # 固有表現認識の実行
        entities = self.ner(contract_text)
        
        # 結果の整理
        result = {
            "status": "success",
            "message": "契約書の処理が完了しました",
            "data": {
                "contract_type": classification[0]["label"],
                "confidence": classification[0]["score"],
                "entities": [
                    {
                        "word": entity["word"],
                        "entity": entity["entity"],
                        "score": entity["score"]
                    }
                    for entity in entities
                ],
                "processed_text": contract_text
            }
        }
        
        return result

def create_gradio_interface():
    """
    Gradioインターフェースを作成する関数
    """
    agent = ContractAgent()
    
    def process_contract_wrapper(text):
        result = agent.process_contract(text)
        
        # 結果の整形
        output = f"""
処理結果:
- 契約タイプ: {result['data']['contract_type']}
- 信頼度: {result['data']['confidence']:.2%}

検出された重要項目:
"""
        for entity in result['data']['entities']:
            output += f"- {entity['word']} ({entity['entity']})\n"
            
        return output
    
    # Gradioインターフェースの作成
    iface = gr.Interface(
        fn=process_contract_wrapper,
        inputs=gr.Textbox(label="契約書テキスト", lines=5),
        outputs=gr.Textbox(label="処理結果", lines=10),
        title="法人契約管理エージェント",
        description="契約書のテキストを入力すると、AIによる分析結果が表示されます。",
        theme="huggingface"
    )
    
    return iface

if __name__ == "__main__":
    # Gradioインターフェースの起動
    iface = create_gradio_interface()
    iface.launch(share=True)