112 lines
2.4 KiB
Python
112 lines
2.4 KiB
Python
|
from enum import Enum
|
||
|
|
from typing import List, Optional
|
||
|
|
|
||
|
|
from pydantic import BaseModel, field_validator
|
||
|
|
|
||
|
|
|
||
|
|
class InvestmentStage(str, Enum):
|
||
|
|
SEED = "SEED"
|
||
|
|
SERIES_A = "SERIES_A"
|
||
|
|
SERIES_B = "SERIES_B"
|
||
|
|
SERIES_C = "SERIES_C"
|
||
|
|
GROWTH = "GROWTH"
|
||
|
|
LATE_STAGE = "LATE_STAGE"
|
||
|
|
|
||
|
|
|
||
|
|
class SectorSchema(BaseModel):
|
||
|
|
id: int
|
||
|
|
name: str
|
||
|
|
|
||
|
|
class Config:
|
||
|
|
from_attributes = True
|
||
|
|
|
||
|
|
|
||
|
|
class InvestorMemberSchema(BaseModel):
|
||
|
|
id: int
|
||
|
|
name: str
|
||
|
|
role: str
|
||
|
|
email: str
|
||
|
|
investor_id: int
|
||
|
|
|
||
|
|
class Config:
|
||
|
|
from_attributes = True
|
||
|
|
|
||
|
|
|
||
|
|
class CompanyMemberSchema(BaseModel):
|
||
|
|
id: int
|
||
|
|
name: Optional[str] = None
|
||
|
|
linkedin: Optional[str] = None
|
||
|
|
role: Optional[str] = None
|
||
|
|
company_id: int
|
||
|
|
|
||
|
|
class Config:
|
||
|
|
from_attributes = True
|
||
|
|
|
||
|
|
|
||
|
|
class CompanySchema(BaseModel):
|
||
|
|
id: int
|
||
|
|
name: str
|
||
|
|
industry: str
|
||
|
|
location: str
|
||
|
|
description: Optional[str] = None # Fixed typo from 'nullabel'
|
||
|
|
founded_year: Optional[int] = None # Changed from str to int to match model
|
||
|
|
website: Optional[str] = None
|
||
|
|
|
||
|
|
@field_validator("founded_year", mode="before")
|
||
|
|
@classmethod
|
||
|
|
def validate_founded_year(cls, v):
|
||
|
|
if v is None or v == "Not Available" or v == "":
|
||
|
|
return None
|
||
|
|
if isinstance(v, str):
|
||
|
|
try:
|
||
|
|
return int(v)
|
||
|
|
except ValueError:
|
||
|
|
return None
|
||
|
|
return v
|
||
|
|
|
||
|
|
class Config:
|
||
|
|
from_attributes = True
|
||
|
|
|
||
|
|
|
||
|
|
class InvestorSchema(BaseModel):
|
||
|
|
id: int
|
||
|
|
name: str
|
||
|
|
description: Optional[str] = None
|
||
|
|
aum: int
|
||
|
|
check_size_lower: int
|
||
|
|
check_size_upper: int
|
||
|
|
geographic_focus: str
|
||
|
|
stage_focus: InvestmentStage
|
||
|
|
number_of_investments: int = 0
|
||
|
|
|
||
|
|
|
||
|
|
class Config:
|
||
|
|
from_attributes = True
|
||
|
|
|
||
|
|
|
||
|
|
class InvestorData(BaseModel):
|
||
|
|
"""Comprehensive investor data schema for LLM processing"""
|
||
|
|
|
||
|
|
investor: InvestorSchema
|
||
|
|
portfolio_companies: List[CompanySchema] = []
|
||
|
|
team_members: List[InvestorMemberSchema] = [] # Changed from TeamMember
|
||
|
|
sectors: List[SectorSchema] = []
|
||
|
|
|
||
|
|
class Config:
|
||
|
|
from_attributes = True
|
||
|
|
|
||
|
|
|
||
|
|
class CompanyData(BaseModel): # Renamed from CompaniesData for consistency
|
||
|
|
company: CompanySchema
|
||
|
|
sectors: List[SectorSchema] = []
|
||
|
|
members: List[CompanyMemberSchema] = [] # Changed to match model relationship name
|
||
|
|
investors: List[InvestorSchema] = []
|
||
|
|
|
||
|
|
class Config:
|
||
|
|
from_attributes = True
|
||
|
|
|
||
|
|
|
||
|
|
class InvestorList(BaseModel):
|
||
|
|
investors: List[InvestorData] = []
|
||
|
|
|