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

cs_models/database/operations.py

@@ -6,7 +6,7 @@ from contextlib import contextmanager
 
 from marshmallow import EXCLUDE
 from sqlalchemy import event
-from sqlalchemy.orm import Session
+from sqlalchemy.orm import sessionmaker, scoped_session
 
 from ..aact_database import engine as engine_aact
 from ..database import engine, get_engine_for_port
@@ -18,6 +18,16 @@ from ..database import engine, get_engine_for_port
 # logging.getLogger('sqlalchemy.engine').setLevel(logging.INFO)
 
 
+# Define a function to create a scoped session with specific configurations
+def create_scoped_session(engine):
+    return scoped_session(
+        sessionmaker(
+            autocommit=False,
+            autoflush=False,
+            bind=engine
+        )
+    )
+
 @contextmanager
 def assert_sql(expected):
     stmts = []
@@ -54,7 +64,8 @@ def session_scope(db: typing.Optional[str] = None):
             bind_engine = get_engine_for_port(tunnel.local_bind_port)
             if db == "aact":
                 bind_engine = engine_aact
-            session = Session(bind=bind_engine)
+            Session = create_scoped_session(bind_engine)
+            session = Session()
 
             try:
                 yield session
@@ -63,12 +74,13 @@ def session_scope(db: typing.Optional[str] = None):
                 session.rollback()
                 raise
             finally:
-                session.close()
+                Session.remove()
     else:
         bind_engine = engine
         if db == "aact":
             bind_engine = engine_aact
-        session = Session(bind=bind_engine)
+        Session = create_scoped_session(bind_engine)
+        session = Session()
 
         try:
             yield session
@@ -77,7 +89,7 @@ def session_scope(db: typing.Optional[str] = None):
             session.rollback()
             raise
         finally:
-            session.close()
+            Session.remove()
 
 
 def create(session, model, schema, obj):

cs_models/resources/ArtifactVectorIndex/models.py

@@ -0,0 +1,32 @@
+from datetime import datetime
+
+from sqlalchemy import (
+    Column,
+    Integer,
+    DateTime,
+    Boolean,
+    String,
+)
+
+from ...database import Base
+
+
+class ArtifactVectorIndexModel(Base):
+    __tablename__ = "artifact_vector_index"
+
+    id = Column(Integer, primary_key=True)
+    artifact_id = Column(
+        String(50),
+        nullable=False,
+        index=True,
+    )
+    embedding_created = Column(Boolean, nullable=True)
+    embedding_s3_key = Column(String(128), nullable=True)
+    vector_indexed = Column(Boolean, nullable=True)
+    updated_at = Column(
+        DateTime,
+        nullable=False,
+        # https://stackoverflow.com/questions/58776476/why-doesnt-freezegun-work-with-sqlalchemy-default-values
+        default=lambda: datetime.utcnow(),
+        onupdate=lambda: datetime.utcnow(),
+    )

cs_models/resources/ArtifactVectorIndex/schemas.py

@@ -0,0 +1,16 @@
+from marshmallow import (
+    Schema,
+    fields,
+    validate,
+)
+
+
+class ArtifactVectorIndexResourceSchema(Schema):
+    not_blank = validate.Length(min=1, error='Field cannot be blank')
+
+    id = fields.Integer(dump_only=True)
+    artifact_id = fields.String(required=True)
+    embedding_created = fields.Boolean(allow_none=True)
+    embedding_s3_key = fields.String(allow_none=True)
+    vector_indexed = fields.Boolean(allow_none=True)
+    updated_at = fields.DateTime()

cs_models/resources/ArtifactVectorIndexQueue/models.py

@@ -0,0 +1,34 @@
+from datetime import datetime
+
+from sqlalchemy import (
+    Column,
+    Integer,
+    DateTime,
+    Boolean,
+    String,
+    Text,
+)
+
+from ...database import Base
+
+
+class ArtifactVectorIndexQueueModel(Base):
+    __tablename__ = "artifact_vector_index_queue"
+
+    id = Column(Integer, primary_key=True)
+    artifact_id = Column(
+        String(50),
+        nullable=False,
+        index=True,
+    )
+    attempts = Column(Integer, nullable=True)
+    submitted = Column(Boolean, nullable=True)
+    submitted_date = Column(DateTime, nullable=True)
+    error = Column(Text, nullable=True)
+    updated_at = Column(
+        DateTime,
+        nullable=False,
+        # https://stackoverflow.com/questions/58776476/why-doesnt-freezegun-work-with-sqlalchemy-default-values
+        default=lambda: datetime.utcnow(),
+        onupdate=lambda: datetime.utcnow(),
+    )

cs_models/resources/ArtifactVectorIndexQueue/schemas.py

@@ -0,0 +1,17 @@
+from marshmallow import (
+    Schema,
+    fields,
+    validate,
+)
+
+
+class ArtifactVectorIndexQueueResourceSchema(Schema):
+    not_blank = validate.Length(min=1, error='Field cannot be blank')
+
+    id = fields.Integer(dump_only=True)
+    artifact_id = fields.String(required=True)
+    attempts = fields.Integer(allow_none=True)
+    submitted = fields.Boolean(allow_none=True)
+    submitted_date = fields.DateTime(allow_none=True)
+    error = fields.String(allow_none=True)
+    updated_at = fields.DateTime()

cs_models/resources/AssistantCommand/models.py

@@ -1,7 +1,7 @@
 import enum
 from datetime import datetime
 
-from sqlalchemy import Column, DateTime, Enum, ForeignKey, Integer, Text
+from sqlalchemy import Column, DateTime, Enum, ForeignKey, Integer, Text, Boolean
 from sqlalchemy.orm import relationship
 
 from ...database import Base
@@ -9,10 +9,21 @@ from ...database import Base
 
 class AssistantCommandTypeEnum(enum.Enum):
     SEARCH = "SEARCH"
+    FINAL_ANSWER = "FINAL_ANSWER"
+    INTERMEDIATE_ANSWER = "INTERMEDIATE_ANSWER"
+    ANSWER = "ANSWER"
+    RELEVANCY = "RELEVANCY"
     EXTRACT = "EXTRACT"
     SUMMARIZE = "SUMMARIZE"
     QUESTION_AND_ANSWER = "QUESTION_AND_ANSWER"
     TABLE_GENERATION = "TABLE_GENERATION"
+    CHART_GENERATION = "CHART_GENERATION"
+    DRUG_LIST = "DRUG_LIST"
+    DEAL_LIST = "DEAL_LIST"
+    LIST = "LIST"
+    SMART_GRID = "SMART_GRID"
+    DYNAMIC_TABLE = "DYNAMIC_TABLE"
+    FOLLOW_UP_QUESTIONS = "FOLLOW_UP_QUESTIONS"
 
 
 class AssistantCommandStatusEnum(enum.Enum):
@@ -37,6 +48,8 @@ class AssistantCommandModel(Base):
     label = Column(Text, nullable=False)
     result = Column(Text)
     llm_result = Column(Text)
+    error = Column(Text, nullable=True)
+    internal_doc = Column(Boolean, nullable=True)
     created_at = Column(DateTime, nullable=False, default=lambda: datetime.utcnow())
     updated_at = Column(
         DateTime,

cs_models/resources/AssistantCommand/schemas.py

@@ -103,6 +103,8 @@ class AssistantCommandResourceSchema(Schema):
     lower_neighbour_ids = fields.Method(serialize="serialize_lower_neighbours")
     result = AssistantCommandResultField(required=True, allow_none=True)
     llm_result = AssistantCommandLLMResultField(allow_none=True)
+    error = fields.String(allow_none=True)
+    internal_doc = fields.Boolean(allow_none=True)
     created_at = fields.DateTime(dump_only=True)
     updated_at = fields.DateTime(dump_only=True)
 

cs_models/resources/AssistantCommandArtifact/models.py

@@ -0,0 +1,33 @@
+from datetime import datetime
+
+from sqlalchemy import (
+    Column,
+    Integer,
+    DateTime,
+    ForeignKey,
+    Float,
+    String,
+    Text,
+)
+
+from ...database import Base
+
+
+class AssistantCommandArtifactModel(Base):
+    __tablename__ = "assistant_command_artifacts"
+
+    id = Column(Integer, primary_key=True)
+    assistant_command_id = Column(
+        Integer,
+        ForeignKey('assistant_commands.id'),
+        nullable=False,
+    )
+    artifact_id = Column(String(64), nullable=False)
+    score = Column(Float, nullable=True)
+    updated_at = Column(
+        DateTime,
+        nullable=False,
+        # https://stackoverflow.com/questions/58776476/why-doesnt-freezegun-work-with-sqlalchemy-default-values
+        default=lambda: datetime.utcnow(),
+        onupdate=lambda: datetime.utcnow(),
+    )

cs_models/resources/AssistantCommandArtifact/schemas.py

@@ -0,0 +1,15 @@
+from marshmallow import (
+    Schema,
+    fields,
+    validate,
+)
+
+
+class AssistantCommandArtifactResourceSchema(Schema):
+    not_blank = validate.Length(min=1, error='Field cannot be blank')
+
+    id = fields.Integer(dump_only=True)
+    assistant_command_id = fields.Integer(required=True)
+    artifact_id = fields.String(required=True)
+    score = fields.Float(allow_none=True)
+    updated_at = fields.DateTime()

cs_models/resources/AssistantCommandChart/models.py

@@ -0,0 +1,37 @@
+from datetime import datetime
+
+from sqlalchemy import (
+    Column,
+    Integer,
+    DateTime,
+    ForeignKey,
+    Boolean,
+    Text,
+)
+
+from ...database import Base
+
+
+class AssistantCommandChartModel(Base):
+    __tablename__ = "assistant_command_charts"
+
+    id = Column(Integer, primary_key=True)
+    assistant_command_id = Column(
+        Integer,
+        ForeignKey('assistant_commands.id'),
+        nullable=False,
+    )
+    file_id = Column(
+        Integer,
+        ForeignKey('files.id'),
+        nullable=False,
+    )
+    chart_info = Column(Text, nullable=True)
+    is_deleted = Column(Boolean, nullable=True)
+    updated_at = Column(
+        DateTime,
+        nullable=False,
+        # https://stackoverflow.com/questions/58776476/why-doesnt-freezegun-work-with-sqlalchemy-default-values
+        default=lambda: datetime.utcnow(),
+        onupdate=lambda: datetime.utcnow(),
+    )

cs_models/resources/AssistantCommandChart/schemas.py

@@ -0,0 +1,16 @@
+from marshmallow import (
+    Schema,
+    fields,
+    validate,
+)
+
+
+class AssistantCommandChartResourceSchema(Schema):
+    not_blank = validate.Length(min=1, error='Field cannot be blank')
+
+    id = fields.Integer(dump_only=True)
+    assistant_command_id = fields.Integer(required=True)
+    file_id = fields.Integer(required=True)
+    chart_info = fields.String(allow_none=True)
+    is_deleted = fields.Boolean(allow_none=True)
+    updated_at = fields.DateTime()

cs_models/resources/AssistantSession/models.py

@@ -1,7 +1,7 @@
 """Models for storing Mindgram assistant sessions."""
 from datetime import datetime
 
-from sqlalchemy import Column, DateTime, Integer, String
+from sqlalchemy import Column, DateTime, Integer, String, Boolean, ForeignKey
 from sqlalchemy.orm import relationship
 
 from ..AssistantUserQuery.models import AssistantUserQueryModel
@@ -17,11 +17,19 @@ class AssistantSessionModel(Base):
     user_id = Column(String(128), nullable=False)
     type = Column(String(50), nullable=False)
     label = Column(String(128), nullable=False)
+    display_label = Column(String(128), nullable=True)
     created_at = Column(
         DateTime,
         nullable=False,
         default=lambda: datetime.utcnow(),
     )
+    is_deleted = Column(Boolean, nullable=True)
+    internal_doc_only = Column(Boolean, nullable=True)
+    workbook_id = Column(
+        Integer,
+        ForeignKey("workbooks.id"),
+        nullable=True,
+    )
     updated_at = Column(
         DateTime,
         nullable=False,

cs_models/resources/AssistantSession/schemas.py

@@ -9,12 +9,16 @@ class AssistantSessionResourceSchema(Schema):
     user_id = fields.String(required=True)
     type = fields.String(required=True)
     label = fields.String(required=True)
+    display_label = fields.String(allow_none=True)
     user_queries = fields.Nested(
         AssistantUserQueryResourceSchema(exclude=("session_id",)),
         many=True,
         dump_only=True,
     )
+    is_deleted = fields.Boolean(allow_none=True)
     created_at = fields.DateTime(dump_only=True)
+    internal_doc_only = fields.Boolean(allow_none=True)
+    workbook_id = fields.Integer(allow_none=True)
     updated_at = fields.DateTime(dump_only=True)
 
     __table_args__ = (UniqueConstraint("user_id", "label"),)

cs_models/resources/AssistantUserQuery/models.py

@@ -1,7 +1,7 @@
 import enum
 from datetime import datetime
 
-from sqlalchemy import Column, DateTime, Enum, ForeignKey, Integer, String
+from sqlalchemy import Column, DateTime, Enum, ForeignKey, Integer, String, Text
 from sqlalchemy.orm import relationship
 
 from ...database import Base
@@ -28,6 +28,10 @@ class AssistantUserQueryModel(Base):
         String,
         nullable=False,
     )
+    filter = Column(
+        Text,
+        nullable=True,
+    )
     type = Column(
         String,
         nullable=False,

cs_models/resources/AssistantUserQuery/schemas.py

@@ -1,15 +1,54 @@
+import json
 from marshmallow import Schema, fields
-
 from .models import AssistantUserQueryStatusEnum
 from ..AssistantCommand.schemas import (
     AssistantCommandResourceSchema,
 )
 
 
+class AssistantUserQueryFilterField(fields.Field):
+    """Field that stores result for the Assistant command."""
+
+    def _serialize(self, value, attr, obj, **kwargs):
+        """
+        In the DB, the `result` field is a text field. We persist
+        data by performing the following:
+
+        AssistantCommandModel(
+            ...
+            result=json.dumps({...}),
+        )
+
+        So here we need to perform the inverse operation (i.e `json.loads(..)`)
+        """
+        if value is None:
+            return None
+        return json.loads(value)
+
+    def _deserialize(self, value, attr, data, **kwargs):
+        """
+        In the DB, the `result` field is a text field. We persist
+        data by performing the following:
+
+
+
+        AssistantCommandModel(
+            ***AssistantCommandResourceSchema().load({
+                ...
+                "result": [{"some_key": 1}, {"some_key": 2}],
+            }),
+        )
+        """
+        if value is None:
+            return None
+        return json.dumps(value)
+
+
 class AssistantUserQueryResourceSchema(Schema):
     id = fields.Integer(dump_only=True)
     session_id = fields.Integer(required=True)
     value = fields.String(required=True)
+    filter = AssistantUserQueryFilterField(allow_none=True)
     type = fields.String(required=True)
     status = fields.Enum(AssistantUserQueryStatusEnum, required=True, by_value=True)
     commands = fields.Nested(

cs_models/resources/B3CRSS/models.py

@@ -0,0 +1,25 @@
+from datetime import datetime
+
+from sqlalchemy import (
+    Column,
+    String,
+    Integer,
+    DateTime,
+)
+
+from ...database import Base
+
+
+class B3CRSSModel(Base):
+    __tablename__ = "b3c_rss"
+
+    id = Column(Integer, primary_key=True)
+    article_id = Column(String(128), nullable=False, index=True)
+    article_date = Column(DateTime, nullable=False)
+    updated_at = Column(
+        DateTime,
+        nullable=False,
+        # https://stackoverflow.com/questions/58776476/why-doesnt-freezegun-work-with-sqlalchemy-default-values
+        default=lambda: datetime.utcnow(),
+        onupdate=lambda: datetime.utcnow(),
+    )

cs_models/resources/B3CRSS/schemas.py

@@ -0,0 +1,14 @@
+from marshmallow import (
+    Schema,
+    fields,
+    validate,
+)
+
+
+class B3CRSSResourceSchema(Schema):
+    not_blank = validate.Length(min=1, error='Field cannot be blank')
+
+    id = fields.Integer(dump_only=True)
+    article_id = fields.String(required=True)
+    article_date = fields.DateTime(required=True)
+    updated_at = fields.DateTime()

cs_models/resources/CampaignTracker/models.py

@@ -0,0 +1,29 @@
+from datetime import datetime
+
+from sqlalchemy import (
+    Column,
+    Integer,
+    DateTime,
+    String,
+    Text,
+)
+
+from ...database import Base
+
+
+class CampaignTrackerModel(Base):
+    __tablename__ = "campaign_tracker"
+
+    id = Column(Integer, primary_key=True)
+    name = Column(String(128), nullable=False)
+    email = Column(String(128), nullable=True)
+    company = Column(String(128), nullable=True)
+    promo_code = Column(String(50), nullable=True)
+    campaign_description = Column(Text, nullable=True)
+    updated_at = Column(
+        DateTime,
+        nullable=False,
+        # https://stackoverflow.com/questions/58776476/why-doesnt-freezegun-work-with-sqlalchemy-default-values
+        default=lambda: datetime.utcnow(),
+        onupdate=lambda: datetime.utcnow(),
+    )

cs_models/resources/CampaignTracker/schemas.py

@@ -0,0 +1,17 @@
+from marshmallow import (
+    Schema,
+    fields,
+    validate,
+)
+
+
+class CampaignTrackerResourceSchema(Schema):
+    not_blank = validate.Length(min=1, error='Field cannot be blank')
+
+    id = fields.Integer(dump_only=True)
+    name = fields.String(required=True)
+    email = fields.String(allow_none=True)
+    company = fields.String(allow_none=True)
+    promo_code = fields.String(allow_none=True)
+    campaign_description = fields.String(allow_none=True)
+    updated_at = fields.DateTime(dump_only=True)

cs_models/resources/CompanyDrugCount/models.py

@@ -0,0 +1,28 @@
+from sqlalchemy import (
+    Column,
+    Integer,
+    String,
+    DateTime,
+)
+from datetime import datetime
+
+from ...database import Base
+
+
+class CompanyDrugCountModel(Base):
+    __tablename__ = 'company_drug_count'
+
+    id = Column(Integer, primary_key=True)
+    company_id = Column(
+        String(50),
+        nullable=False,
+        index=True
+    )
+    drug_count = Column(Integer, nullable=False)
+    updated_at = Column(
+        DateTime,
+        nullable=False,
+        # https://stackoverflow.com/questions/58776476/why-doesnt-freezegun-work-with-sqlalchemy-default-values
+        default=lambda: datetime.utcnow(),
+        onupdate=lambda: datetime.utcnow(),
+    )

cs_models/resources/CompanyDrugCount/schemas.py

@@ -0,0 +1,11 @@
+from marshmallow import (
+    Schema,
+    fields,
+)
+
+
+class CompanyDrugCountResourceSchema(Schema):
+    id = fields.Integer(dump_only=True)
+    company_id = fields.String(required=True)
+    drug_count = fields.Integer(required=True)
+    updated_at = fields.DateTime(dump_only=True)

cs_models/resources/CompanyFiling/models.py

@@ -40,6 +40,7 @@ class CompanyFilingModel(Base):
     is_deleted = Column(Boolean, nullable=True)
     unprocessed = Column(Boolean, nullable=True)
     reviewed = Column(Boolean, nullable=True)
+    page_count = Column(Integer, nullable=True)
     updated_at = Column(
         DateTime,
         nullable=False,

cs_models/resources/CompanyFiling/schemas.py

@@ -24,6 +24,7 @@ class CompanyFilingResourceSchema(Schema):
     is_deleted = fields.Boolean(allow_none=True)
     unprocessed = fields.Boolean(allow_none=True)
     reviewed = fields.Boolean(allow_none=True)
+    page_count = fields.Integer(allow_none=True)
     updated_at = fields.DateTime(dump_only=True)
 
     @pre_load

cs_models/resources/CompanyMarketCap/models.py

@@ -0,0 +1,29 @@
+from sqlalchemy import (
+    Column,
+    Integer,
+    String,
+    DateTime,
+    DECIMAL,
+)
+from datetime import datetime
+
+from ...database import Base
+
+
+class CompanyMarketCapModel(Base):
+    __tablename__ = 'company_market_cap'
+
+    id = Column(Integer, primary_key=True)
+    company_id = Column(
+        String(50),
+        nullable=False,
+        index=True
+    )
+    market_cap = Column(DECIMAL(20, 2), nullable=False)
+    updated_at = Column(
+        DateTime,
+        nullable=False,
+        # https://stackoverflow.com/questions/58776476/why-doesnt-freezegun-work-with-sqlalchemy-default-values
+        default=lambda: datetime.utcnow(),
+        onupdate=lambda: datetime.utcnow(),
+    )

cs_models/resources/CompanyMarketCap/schemas.py

@@ -0,0 +1,11 @@
+from marshmallow import (
+    Schema,
+    fields,
+)
+
+
+class CompanyMarketCapResourceSchema(Schema):
+    id = fields.Integer(dump_only=True)
+    company_id = fields.String(required=True)
+    market_cap = fields.Decimal(required=True)
+    updated_at = fields.DateTime(dump_only=True)

cs_models/resources/CompanyOUS/models.py

@@ -6,6 +6,7 @@ from sqlalchemy import (
     Boolean,
     ForeignKey,
     Text,
+    DECIMAL,
 )
 from datetime import datetime
 
@@ -20,6 +21,7 @@ class CompanyOUSModel(Base):
     name = Column(String(190), unique=True, nullable=False)
     ticker = Column(String(20), nullable=True, index=True)
     exchange = Column(String(100), nullable=True)
+    market_cap = Column(DECIMAL(20, 2), nullable=False)
     company_url = Column(String(255), nullable=True)
     pipeline_url = Column(String(255), nullable=True)
     ir_url = Column(String(255), nullable=True)