feat: use langchain for generic llm usage

Author: somehowchris

requirements.txt

@@ -55,5 +55,12 @@ anthropic
 redis[hiredis]==5.0.3
 uwsgi==2.0.25.1
 psycopg2-binary==2.9.9
+sqlalchemy_easy_softdelete
+sqlalchemy_serializer
+pyruvate
+langchain
+langfuse
+langchain-openai
+langchain-anthropic
 
 -e .

src/riskmatrix/__init__.py

@@ -14,7 +14,8 @@
 from riskmatrix.security_policy import SessionSecurityPolicy
 from openai import OpenAI
 from anthropic import Anthropic
-
+from langchain_openai import ChatOpenAI
+from langchain_anthropic import ChatAnthropic
 
 from typing import TYPE_CHECKING
 if TYPE_CHECKING:
@@ -87,31 +88,46 @@ def main(
             enable_tracing=True,
             send_default_pii=True
         )
-        print("configured sentry")
-        print(sentry_dsn)
 
     with Configurator(settings=settings, root_factory=root_factory) as config:
         includeme(config)
 
         if openai_apikey := settings.get('openai_api_key'):
-
-            openai_client = OpenAI(
-                api_key=openai_apikey
+            openai_client = ChatOpenAI(
+                api_key=openai_apikey,
+                model = "gpt-4o-mini",
+                temperature=0.7
             )
             config.add_request_method(
                 lambda r: openai_client,
-                'openai',
+                'llm',
                 reify=True
             )
-        if anthropic_apikey := settings.get('anthropic_api_key'):
-            anthropic_client = Anthropic(
-                api_key=anthropic_apikey
+        elif anthropic_apikey := settings.get('anthropic_api_key'):
+            anthropic_client = ChatAnthropic(
+                api_key=anthropic_apikey,
+                model="claude-3-5-sonnet-20240620",
+                temperature=0.7
             )
             config.add_request_method(
                 lambda r: anthropic_client,
-                'anthropic',
+                'llm',
+                reify=True
+            )
+            
+        if langfuse_host := settings.get("langfuse_host"):
+            from langfuse.callback import CallbackHandler
+            langfuse_handler = CallbackHandler(
+                secret_key=settings.get("langfuse_secret_key"),
+                public_key=settings.get("langfuse_public_key"),
+                host=langfuse_host,
+            )
+            config.add_request_method(
+                lambda r: langfuse_handler,
+                'langfuse',
                 reify=True
             )
+ 
 
     app = config.make_wsgi_app()
     return Fanstatic(app, versioning=True)

src/riskmatrix/layouts/steps.py

@@ -1,7 +1,11 @@
+from markupsafe import Markup
 from riskmatrix.i18n import _
 
 from typing import NamedTuple
 from typing import TYPE_CHECKING
+
+from riskmatrix.models.risk_assessment import RiskAssessment
+from riskmatrix.models.risk_assessment_info import RiskAssessmentInfo, RiskAssessmentState
 if TYPE_CHECKING:
     from pyramid.interfaces import IRequest
 
@@ -16,6 +20,12 @@ class Step(NamedTuple):
 
 
 def steps(context: 'Organization', request: 'IRequest') -> 'RenderData':
+    assessments = request.dbsession.query(RiskAssessmentInfo).filter(
+        RiskAssessmentInfo.organization_id == context.id,
+        RiskAssessmentInfo.state == RiskAssessmentState.OPEN,
+    ).all()
+
+    t = request.dbsession.query(RiskAssessment).filter(RiskAssessment.risk_assessment_info_id.in_([a.id for a in assessments]), RiskAssessment.likelihood == None, RiskAssessment.impact == None).count()
     return {
         'steps': [
             Step(
@@ -35,10 +45,11 @@ def steps(context: 'Organization', request: 'IRequest') -> 'RenderData':
                 request.route_url('generate_risk_matrix')
             ),
             Step(
-                _('Plan Actions'),
+                Markup('<s>Plan Actions</s>'),
                 '#',
                 disabled=True
             ),
+            Step(_("Finish Assessment"), request.route_url('finish_assessment'), disabled=False),#t > 0 or len(assessments) == 0),
         ]
     }
 
@@ -54,5 +65,6 @@ def show_steps(request: 'IRequest') -> bool:
             'assess_impact',
             'assess_likelihood',
             'generate_risk_matrix',
+            'finish_assessment'
         )
     return False

src/riskmatrix/models/__init__.py

@@ -12,6 +12,7 @@
 from .risk_assessment import RiskAssessment, RiskMatrixAssessment
 from .risk_catalog import RiskCatalog
 from .risk_category import RiskCategory
+from .risk_assessment_info import RiskAssessmentInfo
 from .user import User
 from .password_change_token import PasswordChangeToken
 
@@ -63,5 +64,6 @@ def includeme(config: 'Configurator') -> None:
     'RiskCategory',
     'RiskMatrixAssessment',
     'User',
-    'PasswordChangeToken'
+    'PasswordChangeToken',
+    'RiskAssessmentInfo',
 )

src/riskmatrix/models/asset.py

@@ -1,6 +1,7 @@
 from datetime import datetime
 from pyramid.authorization import Allow
 from sedate import utcnow
+from riskmatrix.orm.softdelete_base import SoftDeleteMixin
 from sqlalchemy import ForeignKey
 from sqlalchemy import UniqueConstraint
 from sqlalchemy.ext.hybrid import hybrid_property
@@ -25,8 +26,9 @@
     from riskmatrix.models import RiskAssessment
     from riskmatrix.types import ACL
 
+from sqlalchemy_serializer import SerializerMixin
 
-class Asset(Base):
+class Asset(Base, SoftDeleteMixin, SerializerMixin):
 
     __tablename__ = 'asset'
     __table_args__ = (
@@ -48,6 +50,7 @@ class Asset(Base):
     assessments: Mapped[list['RiskAssessment']] = relationship(
         back_populates='asset'
     )
+    
     organization: Mapped['Organization'] = relationship(
         back_populates='assets'
     )

src/riskmatrix/models/risk.py

@@ -1,6 +1,7 @@
 from datetime import datetime
 from pyramid.authorization import Allow
 from sedate import utcnow
+from riskmatrix.orm.softdelete_base import SoftDeleteMixin
 from sqlalchemy import ForeignKey
 from sqlalchemy import ForeignKeyConstraint
 from sqlalchemy import UniqueConstraint
@@ -25,8 +26,9 @@
     from riskmatrix.models import RiskCatalog
     from riskmatrix.types import ACL
 
+from sqlalchemy_serializer import SerializerMixin
 
-class Risk(Base):
+class Risk(SoftDeleteMixin, Base, SerializerMixin):
 
     __tablename__ = 'risk'
     __table_args__ = (

src/riskmatrix/models/risk_assessment.py

@@ -1,6 +1,7 @@
 from datetime import datetime
 from pyramid.authorization import Allow
 from sedate import utcnow
+from riskmatrix.orm.softdelete_base import SoftDeleteMixin
 from sqlalchemy import ForeignKey
 from sqlalchemy import UniqueConstraint
 from sqlalchemy.ext.hybrid import hybrid_property
@@ -12,22 +13,25 @@
 
 from riskmatrix.models import Asset
 from riskmatrix.models import Risk
+from riskmatrix.models.risk_assessment_info import RiskAssessmentInfo
 from riskmatrix.orm.meta import Base
 from riskmatrix.orm.meta import UUIDStr
 from riskmatrix.orm.meta import UUIDStrPK
-
+from sqlalchemy.types import JSON
+from dataclasses import dataclass
+from sqlalchemy_serializer import SerializerMixin
 
 from typing import Any, ClassVar
 from typing import TYPE_CHECKING
 if TYPE_CHECKING:
     from riskmatrix.types import ACL
 
 
-class RiskAssessment(Base):
+class RiskAssessment(SoftDeleteMixin, Base, SerializerMixin):
 
     __tablename__ = 'risk_assessment'
     __table_args__ = (
-        UniqueConstraint('risk_id', 'asset_id'),
+        UniqueConstraint('risk_id', 'asset_id', 'risk_assessment_info_id'),
     )
 
     id: Mapped[UUIDStrPK]
@@ -40,33 +44,53 @@ class RiskAssessment(Base):
         index=True,
     )
 
+    risk_assessment_info_id: Mapped[UUIDStr] = mapped_column(
+        ForeignKey('risk_assessment_info.id', ondelete='CASCADE'),
+        index=True,
+    )
+
     meta: Mapped[dict[str, Any]] = mapped_column(default={})
     impact: Mapped[int | None]
     likelihood: Mapped[int | None]
 
     created: Mapped[datetime] = mapped_column(default=utcnow)
     modified: Mapped[datetime | None] = mapped_column(onupdate=utcnow)
 
+    state_at_finish: Mapped[JSON | None] = mapped_column(JSON, nullable=True)
+
     risk: Mapped[Risk] = relationship(
         back_populates='assessments',
         lazy='joined'
     )
+
     asset: Mapped[Asset] = relationship(
         back_populates='assessments',
         lazy='joined'
     )
 
+    risk_assessment_info: Mapped[RiskAssessmentInfo] = relationship(
+        back_populates='assessments',
+        lazy='joined'
+    )
+
+    risk_assessment_info_id: Mapped[UUIDStr] = mapped_column(
+        ForeignKey('risk_assessment_info.id', ondelete='CASCADE'),
+        index=True,
+    )
+
     def __init__(
         self,
         asset:  Asset,
         risk:   Risk,
+        info:  'RiskAssessmentInfo',
         **meta: Any
     ):
         self.id = str(uuid4())
         self.created = utcnow()
         self.asset = asset
         self.risk = risk
         self.meta = meta
+        self.risk_assessment_info = info
 
     @validates('impact', 'likelihood')
     def ensure_larger_than_one(

src/riskmatrix/models/risk_assessment_info.py

@@ -0,0 +1,66 @@
+from dataclasses import dataclass
+from sedate import utcnow
+from riskmatrix.orm.meta import Base
+import enum
+from sqlalchemy import Enum, ForeignKey
+from sqlalchemy.orm import mapped_column
+from sqlalchemy.orm import relationship
+from sqlalchemy.orm import Mapped
+from riskmatrix.orm.meta import UUIDStr
+from riskmatrix.orm.meta import UUIDStrPK
+from sqlalchemy import UniqueConstraint
+from sqlalchemy import Column
+from datetime import datetime
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from riskmatrix.models import RiskAssessment
+    
+
+
+
+class RiskAssessmentState(enum.Enum):
+    OPEN = 0
+    FINISHED = 1
+
+    def __str__(self) -> str:
+        names = {
+            RiskAssessmentState.OPEN: "Offen",
+            RiskAssessmentState.FINISHED: "Geschlossen"
+        }
+        return names[self]
+
+from sqlalchemy_serializer import SerializerMixin
+
+
+class RiskAssessmentInfo(Base, SerializerMixin):
+
+    __tablename__ = 'risk_assessment_info'
+
+    id: Mapped[UUIDStrPK] = mapped_column()
+    organization_id: Mapped[UUIDStr] = mapped_column(
+        ForeignKey('organization.id', ondelete='CASCADE'),
+        index=True,
+    )
+    
+    name: Mapped[str] = mapped_column(nullable=True)
+
+    state: Mapped[RiskAssessmentState] = mapped_column(
+        Enum(RiskAssessmentState),
+        default=RiskAssessmentState.OPEN,
+    )
+
+    created: Mapped[datetime] = mapped_column(default=utcnow)
+    modified: Mapped[datetime | None] = mapped_column(onupdate=utcnow)
+
+    finished_at: Mapped[datetime | None] = mapped_column()
+
+    assessments: Mapped[list['RiskAssessment']] = relationship(
+        back_populates='risk_assessment_info'
+    )
+
+    def __init__(self, organization_id: str):
+        self.organization_id = organization_id
+        self.state = RiskAssessmentState.OPEN
+        self.finished_at = None

src/riskmatrix/models/risk_catalog.py

@@ -1,6 +1,7 @@
 from datetime import datetime
 from pyramid.authorization import Allow
 from sedate import utcnow
+from riskmatrix.orm.softdelete_base import SoftDeleteMixin
 from sqlalchemy import ForeignKey
 from sqlalchemy import UniqueConstraint
 from sqlalchemy.orm import mapped_column
@@ -15,6 +16,7 @@
 from riskmatrix.orm.meta import UUIDStr
 from riskmatrix.orm.meta import UUIDStrPK
 
+from sqlalchemy_serializer import SerializerMixin
 
 from typing import TYPE_CHECKING
 if TYPE_CHECKING:
@@ -23,7 +25,7 @@
     from riskmatrix.types import ACL
 
 
-class RiskCatalog(Base):
+class RiskCatalog(Base, SoftDeleteMixin, SerializerMixin):
 
     __tablename__ = 'risk_catalog'
     __table_args__ = (

src/riskmatrix/models/risk_category.py

@@ -1,6 +1,7 @@
 from datetime import datetime
 from pyramid.authorization import Allow
 from sedate import utcnow
+from riskmatrix.orm.softdelete_base import SoftDeleteMixin
 from sqlalchemy import ForeignKey
 from sqlalchemy import UniqueConstraint
 from sqlalchemy.orm import mapped_column
@@ -22,7 +23,7 @@
     from riskmatrix.types import ACL
 
 
-class RiskCategory(Base):
+class RiskCategory(Base, SoftDeleteMixin):
 
     __tablename__ = 'risk_category'
     __table_args__ = (