cs-models 0.0.827__py3-none-any.whl → 0.0.847__py3-none-any.whl

cs_models/resources/CompanyOUS/models.py

@@ -28,6 +28,8 @@ class CompanyOUSModel(Base):
     is_activated = Column(Boolean, nullable=True)
     is_searchable = Column(Boolean, nullable=True)
     last_crawl_date = Column(DateTime, nullable=True)
+    last_pipeline_crawl_date = Column(DateTime, nullable=True)
+    pipeline_crawl_enabled = Column(Boolean, nullable=True, default=True)
     industry_type = Column(String(50), nullable=True)
     company_sec_link = Column(
         Integer,

cs_models/resources/CompanyOUS/schemas.py

@@ -20,6 +20,8 @@ class CompanyOUSResourceSchema(Schema):
     is_activated = fields.Boolean(allow_none=True)
     is_searchable = fields.Boolean(allow_none=True)
     last_crawl_date = fields.DateTime(allow_none=True)
+    last_pipeline_crawl_date = fields.DateTime(allow_none=True)
+    pipeline_crawl_enabled = fields.Boolean(allow_none=True)
     industry_type = fields.String(allow_none=True)
     company_sec_link = fields.Integer(allow_none=True)
     relevant_links = fields.String(allow_none=True)
@@ -43,6 +45,8 @@ class CompanyOUSUpdateSchema(Schema):
     pipeline_url = fields.String(allow_none=True)
     ir_url = fields.String(allow_none=True)
     last_crawl_date = fields.DateTime(allow_none=True)
+    last_pipeline_crawl_date = fields.DateTime(allow_none=True)
+    pipeline_crawl_enabled = fields.Boolean(allow_none=True)
     industry_type = fields.String(allow_none=True)
     company_sec_link = fields.Integer(allow_none=True)
     updated_at = fields.DateTime(dump_only=True)

cs_models/resources/CompanySEC/models.py

@@ -32,6 +32,8 @@ class CompanySECModel(Base):
     discarded = Column(Boolean, nullable=True)
     skip_sec = Column(Boolean, nullable=True)
     last_crawl_date = Column(DateTime, nullable=True)
+    last_pipeline_crawl_date = Column(DateTime, nullable=True)
+    pipeline_crawl_enabled = Column(Boolean, nullable=True, default=True)
     industry_type = Column(String(50), nullable=True)
     relevant_links = Column(Text, nullable=True)
     notes = Column(Text, nullable=True)

cs_models/resources/CompanySEC/schemas.py

@@ -24,6 +24,8 @@ class CompanySECResourceSchema(Schema):
     discarded = fields.Boolean(allow_none=True)
     skip_sec = fields.Boolean(allow_none=True)
     last_crawl_date = fields.DateTime(allow_none=True)
+    last_pipeline_crawl_date = fields.DateTime(allow_none=True)
+    pipeline_crawl_enabled = fields.Boolean(allow_none=True)
     industry_type = fields.String(allow_none=True)
     relevant_links = fields.String(allow_none=True)
     notes = fields.String(allow_none=True)
@@ -43,5 +45,7 @@ class CompanySECUpdateSchema(Schema):
     discarded = fields.Boolean(allow_none=True)
     skip_sec = fields.Boolean(allow_none=True)
     last_crawl_date = fields.DateTime(allow_none=True)
+    last_pipeline_crawl_date = fields.DateTime(allow_none=True)
+    pipeline_crawl_enabled = fields.Boolean(allow_none=True)
     industry_type = fields.String(allow_none=True)
     updated_at = fields.DateTime(dump_only=True)

cs_models/resources/DeepResearchAgenticUnit/__init__.py

@@ -0,0 +1,14 @@
+"""DeepResearchAgenticUnit - Self-contained research modules for v2 architecture."""
+
+from .models import DeepResearchAgenticUnitModel, UnitStatusEnum
+from .schemas import (
+    DeepResearchAgenticUnitResourceSchema,
+    DeepResearchAgenticUnitCreateSchema,
+)
+
+__all__ = [
+    "DeepResearchAgenticUnitModel",
+    "UnitStatusEnum",
+    "DeepResearchAgenticUnitResourceSchema",
+    "DeepResearchAgenticUnitCreateSchema",
+]

cs_models/resources/DeepResearchAgenticUnit/models.py

@@ -0,0 +1,123 @@
+"""Models for Deep Research Agentic Units - Self-contained research modules."""
+import enum
+from datetime import datetime
+
+from sqlalchemy import (
+    Column,
+    DateTime,
+    Integer,
+    String,
+    Text,
+    Boolean,
+    ForeignKey,
+    Enum,
+    Float,
+)
+from sqlalchemy.orm import relationship
+
+from ...database import Base
+
+
+class UnitStatusEnum(enum.Enum):
+    """Status enum for agentic units."""
+    PENDING = "PENDING"
+    PLANNING = "PLANNING"       # Internal DAG being planned
+    EXECUTING = "EXECUTING"     # Internal tasks running
+    COMPLETED = "COMPLETED"
+    FAILED = "FAILED"
+    SKIPPED = "SKIPPED"
+
+
+class DeepResearchAgenticUnitModel(Base):
+    """
+    Model for Deep Research Agentic Units.
+
+    An Agentic Unit is a self-contained research module that answers
+    ONE coherent objective using any combination of internal tasks
+    (DISCOVERY, INVESTIGATION, SMART_GRID, AGGREGATION).
+
+    Part of the V2 architecture that replaces rigid upfront DAG decomposition
+    with an adaptive Plan -> Execute -> Evaluate loop.
+    """
+
+    __tablename__ = "deep_research_agentic_units"
+
+    id = Column(Integer, primary_key=True)
+
+    # Parent session
+    session_id = Column(
+        Integer,
+        ForeignKey("deep_research_sessions.id"),
+        nullable=False,
+    )
+
+    # Identity
+    unit_id = Column(String(50), nullable=False)  # Logical ID: "A", "B", "C", etc.
+    unit_label = Column(String(255), nullable=True)  # Human-readable label
+    objective = Column(Text, nullable=False)  # What this unit answers
+
+    # DAG Structure (outer DAG - dependencies between units)
+    depends_on = Column(Text, nullable=True)  # JSON: ["A", "B"] - other unit_ids
+    wave_number = Column(Integer, default=0)  # Computed from dependencies
+
+    # Internal Structure (inner DAG - tasks within this unit)
+    internal_dag = Column(Text, nullable=True)  # JSON: List of internal tasks
+    internal_dag_planned = Column(Boolean, default=False)
+
+    # Execution State
+    status = Column(
+        "status",
+        Enum(UnitStatusEnum),
+        default=UnitStatusEnum.PENDING,
+        nullable=False,
+    )
+
+    # Outputs (populated after execution)
+    result_s3_key = Column(String(512), nullable=True)
+    result_summary = Column(Text, nullable=True)  # Concise answer to objective
+    entities_discovered = Column(Text, nullable=True)  # JSON: Entities found (if any)
+    confidence = Column(Float, nullable=True)  # 0-1 confidence in results
+    gaps_identified = Column(Text, nullable=True)  # JSON: What couldn't be answered
+
+    # Metrics
+    tokens_used = Column(Integer, default=0)
+    tasks_completed = Column(Integer, default=0)
+    tasks_failed = Column(Integer, default=0)
+
+    # Context from dependencies (populated before execution)
+    input_context = Column(Text, nullable=True)  # JSON: Summaries from dependency units
+
+    # Timestamps
+    created_at = Column(
+        DateTime,
+        nullable=False,
+        default=lambda: datetime.utcnow(),
+    )
+    started_at = Column(DateTime, nullable=True)
+    completed_at = Column(DateTime, nullable=True)
+    updated_at = Column(
+        DateTime,
+        nullable=False,
+        default=lambda: datetime.utcnow(),
+        onupdate=lambda: datetime.utcnow(),
+    )
+
+    # Error handling
+    error_message = Column(Text, nullable=True)
+    retry_count = Column(Integer, default=0)
+
+    # Soft delete
+    is_deleted = Column(Boolean, nullable=True)
+
+    # ORM Relationships
+    session = relationship(
+        "DeepResearchSessionModel",
+        back_populates="agentic_units",
+    )
+
+    internal_tasks = relationship(
+        "DeepResearchSubTaskModel",
+        primaryjoin="and_(DeepResearchAgenticUnitModel.id==DeepResearchSubTaskModel.agentic_unit_id, "
+                    "or_(DeepResearchSubTaskModel.is_deleted==False, DeepResearchSubTaskModel.is_deleted==None))",
+        back_populates="agentic_unit",
+    )

cs_models/resources/DeepResearchAgenticUnit/schemas.py

@@ -0,0 +1,50 @@
+"""Marshmallow schemas for Deep Research Agentic Units."""
+
+from marshmallow import Schema, fields, EXCLUDE
+
+
+class DeepResearchAgenticUnitResourceSchema(Schema):
+    """Schema for reading agentic unit resources."""
+
+    class Meta:
+        unknown = EXCLUDE
+
+    id = fields.Integer(dump_only=True)
+    session_id = fields.Integer(required=True)
+    unit_id = fields.String(required=True)
+    unit_label = fields.String(allow_none=True)
+    objective = fields.String(required=True)
+    depends_on = fields.String(allow_none=True)  # JSON string
+    wave_number = fields.Integer(allow_none=True)
+    internal_dag = fields.String(allow_none=True)  # JSON string
+    internal_dag_planned = fields.Boolean(allow_none=True)
+    status = fields.String(allow_none=True)
+    result_s3_key = fields.String(allow_none=True)
+    result_summary = fields.String(allow_none=True)
+    entities_discovered = fields.String(allow_none=True)  # JSON string
+    confidence = fields.Float(allow_none=True)
+    gaps_identified = fields.String(allow_none=True)  # JSON string
+    tokens_used = fields.Integer(allow_none=True)
+    tasks_completed = fields.Integer(allow_none=True)
+    tasks_failed = fields.Integer(allow_none=True)
+    input_context = fields.String(allow_none=True)  # JSON string
+    created_at = fields.DateTime(dump_only=True)
+    started_at = fields.DateTime(allow_none=True)
+    completed_at = fields.DateTime(allow_none=True)
+    updated_at = fields.DateTime(dump_only=True)
+    error_message = fields.String(allow_none=True)
+    retry_count = fields.Integer(allow_none=True)
+
+
+class DeepResearchAgenticUnitCreateSchema(Schema):
+    """Schema for creating agentic units."""
+
+    class Meta:
+        unknown = EXCLUDE
+
+    session_id = fields.Integer(required=True)
+    unit_id = fields.String(required=True)
+    unit_label = fields.String(allow_none=True)
+    objective = fields.String(required=True)
+    depends_on = fields.String(allow_none=True)  # JSON string of unit_ids
+    wave_number = fields.Integer(allow_none=True, load_default=0)

cs_models/resources/DeepResearchSession/__init__.py

@@ -0,0 +1,20 @@
+"""Deep Research Session model and schema exports."""
+from .models import (
+    DeepResearchSessionModel,
+    DeepResearchStatusEnum,
+    HITLStatusEnum,
+)
+from .schemas import (
+    DeepResearchSessionResourceSchema,
+    DeepResearchSessionCreateSchema,
+    DeepResearchSessionProgressSchema,
+)
+
+__all__ = [
+    "DeepResearchSessionModel",
+    "DeepResearchStatusEnum",
+    "HITLStatusEnum",
+    "DeepResearchSessionResourceSchema",
+    "DeepResearchSessionCreateSchema",
+    "DeepResearchSessionProgressSchema",
+]

cs_models/resources/DeepResearchSession/models.py

@@ -0,0 +1,170 @@
+"""Models for Deep Research Sessions - Agentic research orchestration."""
+import enum
+from datetime import datetime
+
+from sqlalchemy import (
+    Column,
+    DateTime,
+    Integer,
+    String,
+    Text,
+    Boolean,
+    ForeignKey,
+    Enum,
+    Float,
+)
+from sqlalchemy.orm import relationship
+
+from ...database import Base
+
+
+class DeepResearchStatusEnum(enum.Enum):
+    """Status enum for deep research sessions."""
+    NOT_STARTED = "NOT_STARTED"
+    DECOMPOSING = "DECOMPOSING"
+    WAITING_USER_CLARIFICATION = "WAITING_USER_CLARIFICATION"
+    WAITING_SCOPE_APPROVAL = "WAITING_SCOPE_APPROVAL"
+    RESEARCHING = "RESEARCHING"
+    SYNTHESIZING = "SYNTHESIZING"
+    WAITING_USER_REVIEW = "WAITING_USER_REVIEW"
+    DRILL_DOWN = "DRILL_DOWN"
+    COMPLETED = "COMPLETED"
+    FAILED = "FAILED"
+    CANCELLED = "CANCELLED"
+
+
+class HITLStatusEnum(enum.Enum):
+    """Human-in-the-loop status enum."""
+    NONE = "NONE"
+    PENDING = "PENDING"
+    RESPONDED = "RESPONDED"
+    TIMEOUT = "TIMEOUT"
+
+
+class DeepResearchSessionModel(Base):
+    """
+    Model for Deep Research Sessions.
+
+    Each session represents a complex research task that is decomposed into
+    parallel sub-tasks executed by independent agents. Supports human-in-the-loop
+    checkpoints for clarification, scope approval, and result review.
+    """
+
+    __tablename__ = "deep_research_sessions"
+
+    id = Column(Integer, primary_key=True)
+
+    # User and organization
+    user_id = Column(String(128), nullable=False)
+    org_id = Column(String(128), nullable=True)
+
+    # Link to workbook (optional)
+    workbook_id = Column(
+        Integer,
+        ForeignKey("workbooks.id"),
+        nullable=True,
+    )
+
+    # Original query and configuration
+    original_query = Column(Text, nullable=False)
+    research_config = Column(Text, nullable=True)  # JSON string of config
+
+    # Execution status
+    status = Column(
+        "status",
+        Enum(DeepResearchStatusEnum),
+        default=DeepResearchStatusEnum.NOT_STARTED,
+        nullable=False,
+    )
+    current_phase = Column(String(50), nullable=True)
+
+    # AWS Step Functions tracking
+    step_function_arn = Column(String(512), nullable=True)
+    step_function_execution_arn = Column(String(512), nullable=True)
+
+    # Progress tracking
+    total_subtasks = Column(Integer, nullable=True, default=0)
+    completed_subtasks = Column(Integer, nullable=True, default=0)
+    failed_subtasks = Column(Integer, nullable=True, default=0)
+
+    # Human-in-the-loop state
+    hitl_status = Column(
+        "hitl_status",
+        Enum(HITLStatusEnum),
+        default=HITLStatusEnum.NONE,
+        nullable=False,
+    )
+    hitl_type = Column(String(50), nullable=True)  # clarification, scope_approval, review
+    hitl_questions = Column(Text, nullable=True)  # JSON array of questions
+    hitl_responses = Column(Text, nullable=True)  # JSON object of responses
+    hitl_task_token = Column(String(1024), nullable=True)  # Step Functions callback token
+    hitl_requested_at = Column(DateTime, nullable=True)
+    hitl_responded_at = Column(DateTime, nullable=True)
+
+    # Results storage
+    final_report_s3_key = Column(String(512), nullable=True)
+    executive_summary = Column(Text, nullable=True)
+    smart_grid_id = Column(
+        Integer,
+        ForeignKey("smart_grids.id"),
+        nullable=True,
+    )
+
+    # Metadata
+    total_citations = Column(Integer, nullable=True, default=0)
+    average_confidence = Column(Float, nullable=True)
+    total_documents_analyzed = Column(Integer, nullable=True, default=0)
+
+    # V2 Architecture Fields (Agentic Units)
+    architecture_version = Column(String(10), default="v1")  # "v1" = old DAG, "v2" = agentic units
+
+    # Orchestration state (V2)
+    current_iteration = Column(Integer, default=0)  # Plan-Execute-Evaluate cycle count
+    max_iterations = Column(Integer, default=10)
+    token_budget = Column(Integer, default=100000)
+    tokens_used = Column(Integer, default=0)
+
+    # Unit tracking (V2)
+    total_units = Column(Integer, default=0)
+    completed_units = Column(Integer, default=0)
+
+    # Evaluation state (V2)
+    last_evaluation = Column(Text, nullable=True)  # JSON: Last evaluator output
+    objective_confidence = Column(Float, nullable=True)  # Overall confidence
+
+    # Error handling
+    error_message = Column(Text, nullable=True)
+    retry_count = Column(Integer, nullable=True, default=0)
+
+    # Soft delete and timestamps
+    is_deleted = Column(Boolean, nullable=True)
+    created_at = Column(
+        DateTime,
+        nullable=False,
+        default=lambda: datetime.utcnow(),
+    )
+    updated_at = Column(
+        DateTime,
+        nullable=False,
+        default=lambda: datetime.utcnow(),
+        onupdate=lambda: datetime.utcnow(),
+    )
+    completed_at = Column(DateTime, nullable=True)
+
+    # ORM Relationships
+    subtasks = relationship(
+        "DeepResearchSubTaskModel",
+        primaryjoin="and_(DeepResearchSessionModel.id==DeepResearchSubTaskModel.session_id, "
+                    "or_(DeepResearchSubTaskModel.is_deleted==False, DeepResearchSubTaskModel.is_deleted==None))",
+        order_by="DeepResearchSubTaskModel.sequence_number",
+        back_populates="session",
+    )
+
+    # V2: Agentic Units relationship
+    agentic_units = relationship(
+        "DeepResearchAgenticUnitModel",
+        primaryjoin="and_(DeepResearchSessionModel.id==DeepResearchAgenticUnitModel.session_id, "
+                    "or_(DeepResearchAgenticUnitModel.is_deleted==False, DeepResearchAgenticUnitModel.is_deleted==None))",
+        order_by="DeepResearchAgenticUnitModel.wave_number",
+        back_populates="session",
+    )

cs_models/resources/DeepResearchSession/schemas.py

@@ -0,0 +1,94 @@
+"""Schemas for Deep Research Sessions."""
+from marshmallow import Schema, fields, validate
+
+from .models import DeepResearchStatusEnum, HITLStatusEnum
+
+
+class DeepResearchSessionResourceSchema(Schema):
+    """Schema for DeepResearchSessionModel."""
+
+    not_blank = validate.Length(min=1, error="Field cannot be blank")
+
+    id = fields.Integer(dump_only=True)
+    user_id = fields.String(required=True, validate=not_blank)
+    org_id = fields.String(allow_none=True)
+    workbook_id = fields.Integer(allow_none=True)
+
+    # Query and config
+    original_query = fields.String(required=True)
+    research_config = fields.String(allow_none=True)
+
+    # Status
+    status = fields.Enum(DeepResearchStatusEnum, by_value=True)
+    current_phase = fields.String(allow_none=True)
+
+    # Step Functions
+    step_function_arn = fields.String(allow_none=True)
+    step_function_execution_arn = fields.String(allow_none=True)
+
+    # Progress
+    total_subtasks = fields.Integer(allow_none=True)
+    completed_subtasks = fields.Integer(allow_none=True)
+    failed_subtasks = fields.Integer(allow_none=True)
+
+    # HITL
+    hitl_status = fields.Enum(HITLStatusEnum, by_value=True)
+    hitl_type = fields.String(allow_none=True)
+    hitl_questions = fields.String(allow_none=True)
+    hitl_responses = fields.String(allow_none=True)
+    hitl_task_token = fields.String(allow_none=True)
+    hitl_requested_at = fields.DateTime(allow_none=True)
+    hitl_responded_at = fields.DateTime(allow_none=True)
+
+    # Results
+    final_report_s3_key = fields.String(allow_none=True)
+    executive_summary = fields.String(allow_none=True)
+    smart_grid_id = fields.Integer(allow_none=True)
+
+    # Metadata
+    total_citations = fields.Integer(allow_none=True)
+    average_confidence = fields.Float(allow_none=True)
+    total_documents_analyzed = fields.Integer(allow_none=True)
+
+    # Error handling
+    error_message = fields.String(allow_none=True)
+    retry_count = fields.Integer(allow_none=True)
+
+    # Timestamps
+    is_deleted = fields.Boolean(allow_none=True)
+    created_at = fields.DateTime(dump_only=True)
+    updated_at = fields.DateTime(dump_only=True)
+    completed_at = fields.DateTime(allow_none=True)
+
+    # Nested relationships (dump only)
+    subtasks = fields.Nested(
+        "DeepResearchSubTaskResourceSchema",
+        many=True,
+        dump_only=True,
+        exclude=("session_id",),
+    )
+
+
+class DeepResearchSessionCreateSchema(Schema):
+    """Schema for creating a new Deep Research Session."""
+
+    not_blank = validate.Length(min=1, error="Field cannot be blank")
+
+    original_query = fields.String(required=True, validate=not_blank)
+    workbook_id = fields.Integer(allow_none=True)
+    research_config = fields.Dict(allow_none=True)  # Will be serialized to JSON
+
+
+class DeepResearchSessionProgressSchema(Schema):
+    """Schema for returning session progress."""
+
+    id = fields.Integer()
+    status = fields.Enum(DeepResearchStatusEnum, by_value=True)
+    current_phase = fields.String(allow_none=True)
+    total_subtasks = fields.Integer()
+    completed_subtasks = fields.Integer()
+    failed_subtasks = fields.Integer()
+    hitl_status = fields.Enum(HITLStatusEnum, by_value=True)
+    hitl_type = fields.String(allow_none=True)
+    hitl_questions = fields.String(allow_none=True)
+    error_message = fields.String(allow_none=True)

cs_models/resources/DeepResearchSubTask/__init__.py

@@ -0,0 +1,20 @@
+"""Deep Research SubTask model and schema exports."""
+from .models import (
+    DeepResearchSubTaskModel,
+    SubTaskStatusEnum,
+    SubTaskTypeEnum,
+)
+from .schemas import (
+    DeepResearchSubTaskResourceSchema,
+    DeepResearchSubTaskSpecSchema,
+    DeepResearchSubTaskCreateSchema,
+)
+
+__all__ = [
+    "DeepResearchSubTaskModel",
+    "SubTaskStatusEnum",
+    "SubTaskTypeEnum",
+    "DeepResearchSubTaskResourceSchema",
+    "DeepResearchSubTaskSpecSchema",
+    "DeepResearchSubTaskCreateSchema",
+]