src/models.py

from datetime import datetime
from typing import Optional, List
from pydantic import BaseModel, Field
from sqlalchemy import Column, Integer, String, DateTime, ForeignKey, Numeric, Text, create_engine
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import relationship, sessionmaker

Base = declarative_base()

class Supplier(Base):
    __tablename__ = "suppliers"
    
    id = Column(Integer, primary_key=True, autoincrement=True)
    name = Column(String(255), nullable=False, unique=True)
    contact_info = Column(Text)
    created_at = Column(DateTime, default=datetime.utcnow)
    
    purchases = relationship("Purchase", back_populates="supplier")

class Customer(Base):
    __tablename__ = "customers"
    
    id = Column(Integer, primary_key=True, autoincrement=True)
    name = Column(String(255), nullable=False)
    email = Column(String(255))
    phone = Column(String(50))
    address = Column(Text)
    created_at = Column(DateTime, default=datetime.utcnow)
    
    sales = relationship("Sale", back_populates="customer")

class Product(Base):
    __tablename__ = "products"
    
    id = Column(Integer, primary_key=True, autoincrement=True)
    name = Column(String(255), nullable=False)
    description = Column(Text)
    category = Column(String(100))
    created_at = Column(DateTime, default=datetime.utcnow)
    
    purchases = relationship("Purchase", back_populates="product")
    sales = relationship("Sale", back_populates="product")

class Purchase(Base):
    __tablename__ = "purchases"
    
    id = Column(Integer, primary_key=True, autoincrement=True)
    supplier_id = Column(Integer, ForeignKey("suppliers.id"))
    product_id = Column(Integer, ForeignKey("products.id"))
    quantity = Column(Integer, nullable=False)
    unit_price = Column(Numeric(10, 2), nullable=False)
    total_cost = Column(Numeric(10, 2), nullable=False)
    purchase_date = Column(DateTime, default=datetime.utcnow)
    notes = Column(Text)
    
    supplier = relationship("Supplier", back_populates="purchases")
    product = relationship("Product", back_populates="purchases")

class Sale(Base):
    __tablename__ = "sales"
    
    id = Column(Integer, primary_key=True, autoincrement=True)
    customer_id = Column(Integer, ForeignKey("customers.id"))
    product_id = Column(Integer, ForeignKey("products.id"))
    quantity = Column(Integer, nullable=False)
    unit_price = Column(Numeric(10, 2), nullable=False)
    total_amount = Column(Numeric(10, 2), nullable=False)
    sale_date = Column(DateTime, default=datetime.utcnow)
    notes = Column(Text)
    
    customer = relationship("Customer", back_populates="sales")
    product = relationship("Product", back_populates="sales")

# Pydantic models for API
class EntityExtraction(BaseModel):
    product: Optional[str] = None
    quantity: Optional[int] = None
    unit: Optional[str] = None  # e.g., "tons", "pieces", "kg"
    supplier: Optional[str] = None
    customer: Optional[str] = None
    unit_price: Optional[float] = None
    total_amount: Optional[float] = None
    transaction_type: str = Field(..., description="'purchase' or 'sale'")
    notes: Optional[str] = None

class ChatbotRequest(BaseModel):
    message: str
    session_id: Optional[str] = None

class PendingTransaction(BaseModel):
    entities: EntityExtraction
    missing_fields: List[str]
    session_id: str
    original_message: str
    clarification_responses: List[str] = []

class ChatbotResponse(BaseModel):
    response: str
    sql_executed: Optional[str] = None
    entities_extracted: Optional[EntityExtraction] = None
    vector_stored: bool = False
    intent_detected: Optional[str] = None
    intent_confidence: Optional[float] = None
    awaiting_clarification: bool = False