koreshield 0.1.5__py3-none-any.whl → 0.2.0__py3-none-any.whl

koreshield_sdk/integrations/langchain.py

@@ -7,13 +7,7 @@ from langchain_core.messages import BaseMessage
 
 from ..client import KoreShieldClient
 from ..async_client import AsyncKoreShieldClient
-from ..types import (
-    DetectionResult,
-    ThreatLevel,
-    RAGDocument,
-    RAGScanResponse,
-    RAGScanConfig,
-)
+from ..types import DetectionResult, ThreatLevel
 from ..exceptions import KoreShieldError
 
 
@@ -278,193 +272,4 @@ def create_async_koreshield_callback(
         block_on_threat=block_on_threat,
         threat_threshold=threat_threshold,
         **kwargs
-    )
-
-
-# RAG Document Scanning Support
-
-class SecureRetriever:
-    """Wrapper for LangChain retrievers that adds automatic RAG security scanning.
-    
-    This class wraps any LangChain retriever and automatically scans retrieved
-    documents for indirect prompt injection attacks before returning them.
-    
-    Example:
-        ```python
-        from langchain.vectorstores import Chroma
-        from koreshield_sdk.integrations.langchain import SecureRetriever
-        
-        # Original retriever
-        base_retriever = vectorstore.as_retriever()
-        
-        # Wrap with security
-        secure_retriever = SecureRetriever(
-            retriever=base_retriever,
-            koreshield_api_key="your-key",
-            block_threats=True,
-            min_confidence=0.3
-        )
-        
-        # Use as normal - automatic scanning
-        docs = secure_retriever.get_relevant_documents("user query")
-        # Threatening documents are automatically filtered
-        ```
-    """
-    
-    def __init__(
-        self,
-        retriever: Any,
-        koreshield_api_key: str,
-        koreshield_base_url: str = "http://localhost:8000",
-        block_threats: bool = True,
-        min_confidence: float = 0.3,
-        enable_cross_document_analysis: bool = True,
-        log_threats: bool = True,
-    ):
-        """Initialize secure retriever.
-        
-        Args:
-            retriever: Base LangChain retriever to wrap
-            koreshield_api_key: KoreShield API key
-            koreshield_base_url: API base URL
-            block_threats: Whether to filter threatening documents
-            min_confidence: Threat confidence threshold (0.0-1.0)
-            enable_cross_document_analysis: Enable multi-doc threat detection
-            log_threats: Log detected threats
-        """
-        self.retriever = retriever
-        self.koreshield = KoreShieldClient(
-            api_key=koreshield_api_key,
-            base_url=koreshield_base_url
-        )
-        self.block_threats = block_threats
-        self.min_confidence = min_confidence
-        self.enable_cross_document_analysis = enable_cross_document_analysis
-        self.log_threats = log_threats
-        
-        # Statistics
-        self.total_scans = 0
-        self.total_threats_detected = 0
-        self.total_documents_blocked = 0
-    
-    def get_relevant_documents(self, query: str) -> List[Any]:
-        """Retrieve and scan documents.
-        
-        Args:
-            query: User's query
-            
-        Returns:
-            List of LangChain documents (threats filtered if enabled)
-        """
-        # Retrieve documents
-        documents = self.retriever.get_relevant_documents(query)
-        
-        if not documents:
-            return documents
-        
-        # Convert to RAG documents
-        rag_documents = []
-        for idx, doc in enumerate(documents):
-            rag_doc = RAGDocument(
-                id=doc.metadata.get("id", f"doc_{idx}"),
-                content=doc.page_content,
-                metadata=doc.metadata
-            )
-            rag_documents.append(rag_doc)
-        
-        # Scan with KoreShield
-        config = RAGScanConfig(
-            min_confidence=self.min_confidence,
-            enable_cross_document_analysis=self.enable_cross_document_analysis
-        )
-        
-        result = self.koreshield.scan_rag_context(
-            user_query=query,
-            documents=rag_documents,
-            config=config
-        )
-        
-        self.total_scans += 1
-        
-        # Handle threats
-        if not result.is_safe:
-            self.total_threats_detected += 1
-            
-            if self.log_threats:
-                print(f"[KoreShield] RAG threat detected: {result.overall_severity}")
-                print(f"[KoreShield] Confidence: {result.overall_confidence:.2f}")
-                print(f"[KoreShield] Vectors: {result.taxonomy.injection_vectors}")
-            
-            if self.block_threats:
-                # Filter out threatening documents
-                safe_rag_docs = result.get_safe_documents(rag_documents)
-                safe_ids = {doc.id for doc in safe_rag_docs}
-                
-                filtered_docs = [
-                    doc for idx, doc in enumerate(documents)
-                    if rag_documents[idx].id in safe_ids
-                ]
-                
-                blocked_count = len(documents) - len(filtered_docs)
-                self.total_documents_blocked += blocked_count
-                
-                if self.log_threats:
-                    print(f"[KoreShield] Filtered {blocked_count} threatening documents")
-                
-                return filtered_docs
-        
-        return documents
-    
-    def get_stats(self) -> Dict[str, Any]:
-        """Get retriever statistics.
-        
-        Returns:
-            Dictionary with scan statistics
-        """
-        return {
-            "total_scans": self.total_scans,
-            "total_threats_detected": self.total_threats_detected,
-            "total_documents_blocked": self.total_documents_blocked,
-            "threat_detection_rate": (
-                self.total_threats_detected / self.total_scans
-                if self.total_scans > 0 else 0.0
-            )
-        }
-
-
-def secure_retriever(
-    retriever: Any,
-    api_key: str,
-    base_url: str = "http://localhost:8000",
-    **kwargs
-) -> SecureRetriever:
-    """Create a secure retriever from any LangChain retriever.
-    
-    Args:
-        retriever: Base LangChain retriever
-        api_key: KoreShield API key
-        base_url: KoreShield API base URL
-        **kwargs: Additional SecureRetriever arguments
-        
-    Returns:
-        SecureRetriever instance
-    
-    Example:
-        ```python
-        from koreshield_sdk.integrations.langchain import secure_retriever
-        
-        safe_retriever = secure_retriever(
-            vectorstore.as_retriever(),
-            api_key="your-key",
-            block_threats=True
-        )
-        
-        docs = safe_retriever.get_relevant_documents("user query")
-        ```
-    """
-    return SecureRetriever(
-        retriever,
-        koreshield_api_key=api_key,
-        koreshield_base_url=base_url,
-        **kwargs
     )

koreshield_sdk/types.py

@@ -88,161 +88,55 @@ class BatchScanResponse(BaseModel):
     processing_time_ms: float
     request_id: str
     timestamp: str
+    version: str
 
 
-# RAG Detection Types
-
-class InjectionVector(str, Enum):
-    """RAG injection vector taxonomy."""
-    EMAIL = "email"
-    DOCUMENT = "document"
-    WEB_SCRAPING = "web_scraping"
-    DATABASE = "database"
-    CHAT_MESSAGE = "chat_message"
-    CUSTOMER_SUPPORT = "customer_support"
-    KNOWLEDGE_BASE = "knowledge_base"
-    API_INTEGRATION = "api_integration"
-    UNKNOWN = "unknown"
-
-
-class OperationalTarget(str, Enum):
-    """RAG operational target taxonomy."""
-    DATA_EXFILTRATION = "data_exfiltration"
-    PRIVILEGE_ESCALATION = "privilege_escalation"
-    ACCESS_CONTROL_BYPASS = "access_control_bypass"
-    CONTEXT_POISONING = "context_poisoning"
-    SYSTEM_PROMPT_LEAKING = "system_prompt_leaking"
-    MISINFORMATION = "misinformation"
-    RECONNAISSANCE = "reconnaissance"
-    UNKNOWN = "unknown"
-
-
-class PersistenceMechanism(str, Enum):
-    """RAG persistence mechanism taxonomy."""
-    SINGLE_TURN = "single_turn"
-    MULTI_TURN = "multi_turn"
-    CONTEXT_PERSISTENCE = "context_persistence"
-    NON_PERSISTENT = "non_persistent"
-
-
-class EnterpriseContext(str, Enum):
-    """Enterprise context taxonomy."""
-    CRM = "crm"
-    SALES = "sales"
-    CUSTOMER_SUPPORT = "customer_support"
-    MARKETING = "marketing"
-    HEALTHCARE = "healthcare"
-    FINANCIAL_SERVICES = "financial_services"
-    GENERAL = "general"
-
-
-class DetectionComplexity(str, Enum):
-    """Detection complexity taxonomy."""
-    LOW = "low"
-    MEDIUM = "medium"
-    HIGH = "high"
-
-
-class RAGDocument(BaseModel):
-    """Document to be scanned in RAG context."""
-    id: str
+class StreamingScanRequest(BaseModel):
+    """Request for streaming security scanning."""
     content: str
-    metadata: Optional[Dict[str, Any]] = Field(default_factory=dict)
+    chunk_size: int = 1000
+    overlap: int = 100
+    context: Optional[Dict[str, Any]] = None
+    user_id: Optional[str] = None
+    session_id: Optional[str] = None
+    metadata: Optional[Dict[str, Any]] = None
 
     model_config = ConfigDict(extra="allow")
 
 
-class DocumentThreat(BaseModel):
-    """Individual document-level threat."""
-    document_id: str
-    severity: ThreatLevel
-    confidence: float
-    patterns_matched: List[str]
-    injection_vectors: List[InjectionVector]
-    operational_targets: List[OperationalTarget]
-    metadata: Optional[Dict[str, Any]] = None
+class StreamingScanResponse(BaseModel):
+    """Response from streaming security scanning."""
+    chunk_results: List[DetectionResult]
+    overall_result: DetectionResult
+    total_chunks: int
+    processing_time_ms: float
+    request_id: str
+    timestamp: str
+    version: str
 
 
-class CrossDocumentThreat(BaseModel):
-    """Cross-document threat detected across multiple documents."""
-    threat_type: str  # "staged_attack", "coordinated_instructions", "temporal_chain"
-    severity: ThreatLevel
-    confidence: float
-    document_ids: List[str]
-    description: str
-    patterns: List[str]
+class SecurityPolicy(BaseModel):
+    """Custom security policy configuration."""
+    name: str
+    description: Optional[str] = None
+    threat_threshold: ThreatLevel = ThreatLevel.MEDIUM
+    blocked_detection_types: List[DetectionType] = Field(default_factory=list)
+    custom_rules: List[Dict[str, Any]] = Field(default_factory=list)
+    allowlist_patterns: List[str] = Field(default_factory=list)
+    blocklist_patterns: List[str] = Field(default_factory=list)
     metadata: Optional[Dict[str, Any]] = None
 
 
-class TaxonomyClassification(BaseModel):
-    """5-dimensional taxonomy classification."""
-    injection_vectors: List[InjectionVector]
-    operational_targets: List[OperationalTarget]
-    persistence_mechanisms: List[PersistenceMechanism]
-    enterprise_contexts: List[EnterpriseContext]
-    detection_complexity: DetectionComplexity
-
-
-class ContextAnalysis(BaseModel):
-    """RAG context analysis results."""
-    document_threats: List[DocumentThreat]
-    cross_document_threats: List[CrossDocumentThreat]
-    statistics: Dict[str, Any]
-
-
-class RAGScanResponse(BaseModel):
-    """Response from RAG context scanning."""
-    is_safe: bool
-    overall_severity: ThreatLevel
-    overall_confidence: float
-    taxonomy: TaxonomyClassification
-    context_analysis: ContextAnalysis
-    request_id: Optional[str] = None
-    timestamp: Optional[str] = None
-
-    def get_threat_document_ids(self) -> List[str]:
-        """Get list of document IDs with detected threats.
-        
-        Returns:
-            List of document IDs that contain threats
-        """
-        threat_ids = set()
-        
-        # From document-level threats
-        for threat in self.context_analysis.document_threats:
-            threat_ids.add(threat.document_id)
-        
-        # From cross-document threats
-        for threat in self.context_analysis.cross_document_threats:
-            threat_ids.update(threat.document_ids)
-        
-        return list(threat_ids)
-    
-    def get_safe_documents(self, original_documents: List[RAGDocument]) -> List[RAGDocument]:
-        """Filter out threatening documents.
-        
-        Args:
-            original_documents: Original list of documents scanned
-            
-        Returns:
-            List of documents without detected threats
-        """
-        threat_ids = set(self.get_threat_document_ids())
-        return [doc for doc in original_documents if doc.id not in threat_ids]
-    
-    def has_critical_threats(self) -> bool:
-        """Check if critical threats were detected.
-        
-        Returns:
-            True if any critical severity threats found
-        """
-        return self.overall_severity == ThreatLevel.CRITICAL
-
-
-class RAGScanRequest(BaseModel):
-    """Request for RAG context scanning"""
-    user_query: str
-    documents: List[RAGDocument]
-    config: Optional[Dict[str, Any]] = Field(default_factory=dict)
-
-    model_config = ConfigDict(extra="allow")
+class PerformanceMetrics(BaseModel):
+    """SDK performance and usage metrics."""
+    total_requests: int = 0
+    total_processing_time_ms: float = 0.0
+    average_response_time_ms: float = 0.0
+    requests_per_second: float = 0.0
+    error_count: int = 0
+    cache_hit_rate: float = 0.0
+    batch_efficiency: float = 0.0
+    streaming_chunks_processed: int = 0
+    uptime_seconds: float = 0.0
+    memory_usage_mb: Optional[float] = None
+    custom_metrics: Dict[str, Any] = Field(default_factory=dict)

koreshield-0.1.5.dist-info/RECORD

@@ -1,13 +0,0 @@
-koreshield-0.1.5.dist-info/licenses/LICENSE,sha256=k3qeCwQxhbOO1GtxA10Do4-_veQzgflqjOp5uZD5mug,1071
-koreshield_sdk/__init__.py,sha256=JXErgUsoxTgM4EU--Os4ZTobARKWj1Mfurln-hNgCQw,785
-koreshield_sdk/async_client.py,sha256=WF4MQVefUJs-YpjVE4qkrP5P9vT6wb5qFJdsdebtOtc,14877
-koreshield_sdk/client.py,sha256=LHuCrHwugzDeoMY5bxmYRmIyRUwJUNgL_Vv3f5ncqpE,13217
-koreshield_sdk/exceptions.py,sha256=3j1FR4VFbe1Vv4i0bofBgQ_ZGwBfpOInBd9OyNQFUxo,945
-koreshield_sdk/py.typed,sha256=8ZJUsxZiuOy1oJeVhsTWQhTG_6pTVHVXk5hJL79ebTk,25
-koreshield_sdk/types.py,sha256=UabFBswT4ckPt2Umwl9FqOBSpPl6RN4FWJPl5qDn5cc,7034
-koreshield_sdk/integrations/__init__.py,sha256=po_sLSND55Wdu1vDmx4Nrjm072HLf04yxmtWj43yv7Y,382
-koreshield_sdk/integrations/langchain.py,sha256=w3BXs3tVk7R4ldFPhAm7qXbJPsHoamY3z2Ke0WPBVas,16542
-koreshield-0.1.5.dist-info/METADATA,sha256=XqNTIRL56qucFtHk2U0l7sfvtWSmCfGMiPwUTslFQ6A,15408
-koreshield-0.1.5.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
-koreshield-0.1.5.dist-info/top_level.txt,sha256=ePw2ZI3SrHZ5CaTRCyj3aya3j_qTcmRAQjoU7s3gAdM,15
-koreshield-0.1.5.dist-info/RECORD,,