backend/agents/profiler.py

import re
from schemas import AgentState, CompanyProfile
from tools.company_search import fetch_regon_data


def extract_nip(text: str) -> str:
    clean_text = text.replace("-", "").replace(" ", "")
    match = re.search(r"\d{10}", clean_text)
    if match:
        return match.group(0)
    return "1234567890"


def calculate_company_size(revenue: float, employment: int) -> str:
    if employment < 10 and revenue <= 2000000:
        return "Mikro"
    return "MŚP"


def profiler_node(state: AgentState):
    if not state.messages:
        return {"current_agent": "supervisor"}

    user_msg = (
        state.messages[-1]["content"]
        if isinstance(state.messages[-1], dict)
        else getattr(state.messages[-1], "content", "")
    )
    nip = extract_nip(user_msg)
    raw_data = fetch_regon_data(nip)

    profile = CompanyProfile(
        nip=nip,
        pkd_codes=raw_data.get("pkd", []),
        voivodeship=raw_data.get("voivodeship", "Mazowieckie"),
        size=calculate_company_size(
            raw_data.get("revenue", 0), raw_data.get("employment", 0)
        ),
    )

    voivodeship_str = (
        f" z województwa {profile.voivodeship}"
        if profile.voivodeship and profile.voivodeship != "Nieznane"
        else ""
    )
    return {
        "profile": profile,
        "messages": [
            {
                "role": "assistant",
                "content": f"Zidentyfikowałem firmę{voivodeship_str}. Jaka jest główna potrzeba inwestycyjna?",
            }
        ],
        "current_agent": "supervisor",
    }