knowledge2 0.4.0__py3-none-any.whl

sdk/models/onboarding.py

@@ -0,0 +1,323 @@
+"""Pydantic models for onboarding types."""
+
+from __future__ import annotations
+
+from sdk.models._base import K2BaseModel
+
+# =============================================================================
+# Gold Labels
+# =============================================================================
+
+
+class GoldLabelDocumentReferenceModel(K2BaseModel):
+    type: str | None = None
+    value: str | None = None
+
+
+class GoldLabelChunkReferenceModel(K2BaseModel):
+    type: str | None = None
+    value: str | int | dict | None = None
+
+
+class GoldLabelEntryModel(K2BaseModel):
+    query: str | None = None
+    document_reference: GoldLabelDocumentReferenceModel | None = None
+    chunk_reference: GoldLabelChunkReferenceModel | None = None
+    relevant_text: str | None = None
+    metadata: dict | None = None
+
+
+class GoldLabelsUploadRequestModel(K2BaseModel):
+    labels: list[GoldLabelEntryModel] | None = None
+    description: str | None = None
+
+
+class ResolvedLabelInfoModel(K2BaseModel):
+    label_id: str
+    query: str
+    chunk_id: str | None
+    document_id: str | None
+    confidence: float
+    resolution_method: str
+
+
+class GoldLabelsUploadResponseModel(K2BaseModel):
+    corpus_id: str
+    total_uploaded: int
+    resolved_count: int
+    unmatched_count: int
+    labels: list[ResolvedLabelInfoModel]
+
+
+class GoldLabelListItemModel(K2BaseModel):
+    id: str
+    query: str
+    chunk_id: str | None
+    document_id: str | None
+    source: str | None
+    confidence: float | None
+    created_at: str
+
+
+class GoldLabelsListResponseModel(K2BaseModel):
+    corpus_id: str
+    total: int
+    labels: list[GoldLabelListItemModel]
+
+
+# =============================================================================
+# Dataset Analysis
+# =============================================================================
+
+
+class DatasetAnalysisRequestModel(K2BaseModel):
+    description: str | None = None
+    auto_bootstrap: bool | None = None
+    bootstrap_num_samples: int | None = None
+    queries_per_chunk: int | None = None
+
+
+class DatasetAnalysisResponseModel(K2BaseModel):
+    analysis_id: str
+    corpus_id: str
+    status: str
+    job_id: str | None
+    estimated_duration_seconds: int | None
+    created_at: str
+
+
+class DatasetAnalysisSummaryModel(K2BaseModel):
+    analysis_id: str
+    status: str
+    current_stage: str | None
+    stages_completed: list[str]
+    has_prompt: bool
+    has_evaluation: bool
+    bootstrap_enabled: bool
+    bootstrap_labels_count: int | None
+    documents_at_analysis: int | None
+    started_at: str
+    created_at: str
+    completed_at: str | None
+    error_message: str | None
+    domain: str | None
+    lexical_strategy: str | None
+    quality_score: float | None
+
+
+class OnboardingStatusResponseModel(K2BaseModel):
+    corpus_id: str
+    latest_analysis: DatasetAnalysisSummaryModel | None
+    gold_labels_count: int
+    synthetic_batches_count: int
+    has_summaries: bool
+    analysis_stale: bool
+
+
+class SearcherPersonaModel(K2BaseModel):
+    who: str
+    goal: str
+    knowledge_state: str
+
+
+class LexicalStrategyModel(K2BaseModel):
+    recommendation: str
+    rationale: str
+    techniques: list[str]
+
+
+class DatasetAnalysisDetailsModel(K2BaseModel):
+    analysis_id: str
+    corpus_id: str
+    status: str
+    created_at: str
+    completed_at: str | None
+    config: dict
+    schema_analysis: dict | None
+    domain: str | None
+    expertise_level: str | None
+    data_relationship: str | None
+    searcher_persona: SearcherPersonaModel | None
+    lexical_strategy: LexicalStrategyModel | None
+    stage1_analysis: dict | None
+    artifact_uri: str | None
+    prompt_uri: str | None
+    bootstrap_enabled: bool
+    bootstrap_labels_count: int | None
+    error_message: str | None
+
+
+# =============================================================================
+# Synthetic Query Generation
+# =============================================================================
+
+
+class SyntheticQueryGenerationRequestModel(K2BaseModel):
+    analysis_id: str | None = None
+    sample_size: int | None = None
+    queries_per_chunk: int | None = None
+    use_document_context: bool | None = None
+    eval_sample_size: int | None = None
+
+
+class SyntheticQueryBatchResponseModel(K2BaseModel):
+    batch_id: str
+    corpus_id: str
+    analysis_id: str
+    status: str
+    job_id: str | None
+    sample_size: int | None
+    queries_per_chunk: int
+    estimated_queries: int | None
+    created_at: str
+
+
+class SyntheticQueryBatchSummaryModel(K2BaseModel):
+    batch_id: str
+    status: str
+    sample_size: int | None
+    queries_per_chunk: int
+    total_chunks_processed: int | None
+    total_queries_generated: int | None
+    created_at: str
+    completed_at: str | None
+
+
+class SyntheticQueryBatchListResponseModel(K2BaseModel):
+    corpus_id: str
+    batches: list[SyntheticQueryBatchSummaryModel]
+
+
+class SyntheticQuerySampleModel(K2BaseModel):
+    chunk_id: str
+    chunk_text_preview: str
+    queries: list[str]
+
+
+class SyntheticQueryBatchDetailsModel(K2BaseModel):
+    batch_id: str
+    corpus_id: str
+    analysis_id: str
+    status: str
+    created_at: str
+    completed_at: str | None
+    sample_size: int | None
+    queries_per_chunk: int
+    use_document_context: bool
+    config: dict
+    total_chunks_processed: int | None
+    total_queries_generated: int | None
+    stats: dict
+    artifact_uri: str | None
+    sample_queries: list[SyntheticQuerySampleModel]
+    error_message: str | None
+
+
+# =============================================================================
+# Evaluation
+# =============================================================================
+
+
+class EvaluationRequestModel(K2BaseModel):
+    batch_id: str | None = None
+    sample_size: int | None = None
+    generate_report: bool | None = None
+    report_formats: list[str] | None = None
+
+
+class EvaluationResponseModel(K2BaseModel):
+    eval_id: str
+    corpus_id: str
+    batch_id: str
+    status: str
+    job_id: str | None
+    created_at: str
+
+
+class EvaluationMetricsModel(K2BaseModel):
+    total_evaluated: int | None = None
+    avg_relevance: float | None = None
+    avg_groundedness: float | None = None
+    avg_style: float | None = None
+    avg_lexical_diversity: float | None = None
+    pass_rate: float | None = None
+    score_distribution: dict | None = None
+
+
+class EvaluationSummaryModel(K2BaseModel):
+    eval_id: str
+    batch_id: str
+    status: str
+    sample_size: int | None
+    metrics: EvaluationMetricsModel | None
+    created_at: str
+    completed_at: str | None
+
+
+class EvaluationListResponseModel(K2BaseModel):
+    corpus_id: str
+    evaluations: list[EvaluationSummaryModel]
+
+
+class EvaluationDetailsModel(K2BaseModel):
+    eval_id: str
+    corpus_id: str
+    batch_id: str
+    status: str
+    created_at: str
+    completed_at: str | None
+    sample_size: int | None
+    config: dict
+    metrics: EvaluationMetricsModel | None
+    artifact_uri: str | None
+    report_uri: str | None
+    sample_results: list[dict]
+    error_message: str | None
+
+
+class EvaluationReportResponseModel(K2BaseModel):
+    eval_id: str
+    corpus_id: str
+    format: str
+    report_uri: str | None
+    metrics: EvaluationMetricsModel | None
+    recommendations: list[str]
+
+
+# =============================================================================
+# Document Summarization
+# =============================================================================
+
+
+class SummarizationRequestModel(K2BaseModel):
+    force_regenerate: bool | None = None
+
+
+class SummarizationResponseModel(K2BaseModel):
+    corpus_id: str
+    job_id: str | None
+    status: str
+    total_documents: int
+    documents_with_summaries: int
+    documents_to_summarize: int
+    created_at: str
+
+
+class SummarizationStatusResponseModel(K2BaseModel):
+    corpus_id: str
+    total_documents: int
+    documents_with_summaries: int
+    coverage_percent: float
+    latest_job_status: str | None
+    latest_job_completed_at: str | None
+
+
+class DocumentSummaryResponseModel(K2BaseModel):
+    document_id: str
+    corpus_id: str
+    summary: str
+    document_type: str | None
+    entities: dict
+    key_facts: list[str]
+    generation_model: str | None
+    created_at: str

sdk/models/orgs.py

@@ -0,0 +1,11 @@
+"""Pydantic models for org types."""
+
+from __future__ import annotations
+
+from sdk.models._base import K2BaseModel
+
+
+class OrgResponseModel(K2BaseModel):
+    id: str | None = None
+    name: str | None = None
+    contact_email: str | None = None

sdk/models/pipelines.py

@@ -0,0 +1,147 @@
+"""Pydantic models for pipeline types."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from sdk.models._base import K2BaseModel
+
+
+class ValidationIssueModel(K2BaseModel):
+    severity: str | None = None
+    code: str | None = None
+    message: str | None = None
+    path: str | None = None
+
+
+class PipelineSpecResponseModel(K2BaseModel):
+    id: str | None = None
+    org_id: str | None = None
+    project_id: str | None = None
+    name: str | None = None
+    description: str | None = None
+    topology: dict[str, Any] | None = None
+    status: str | None = None
+    has_draft: bool | None = None
+    parent_pipeline_spec_id: str | None = None
+    created_by: str | None = None
+    bound_entity_count: int | None = None
+    binding_summary: dict[str, Any] | None = None
+    created_at: str | None = None
+    updated_at: str | None = None
+
+
+class PipelineSpecListResponseModel(K2BaseModel):
+    pipeline_specs: list[PipelineSpecResponseModel]
+    total: int
+    limit: int
+    offset: int
+
+
+class DryRunResultModel(K2BaseModel):
+    valid: bool | None = None
+    issues: list[ValidationIssueModel] | None = None
+    would_create_agents: int | None = None
+    would_create_feeds: int | None = None
+    would_create_corpora: int | None = None
+    would_bind_existing: int | None = None
+
+
+class ApplyResultModel(K2BaseModel):
+    success: bool | None = None
+    created_agent_ids: list[str] | None = None
+    created_feed_ids: list[str] | None = None
+    created_corpus_ids: list[str] | None = None
+    created_subscription_ids: list[str] | None = None
+    bound_entity_ids: list[str] | None = None
+    issues: list[ValidationIssueModel] | None = None
+
+
+class SkippedEntityModel(K2BaseModel):
+    entity_type: str | None = None
+    entity_id: str | None = None
+    entity_name: str | None = None
+    reason: str | None = None
+    other_pipeline_ids: list[str] | None = None
+
+
+class ArchiveResultModel(K2BaseModel):
+    pipeline_spec: PipelineSpecResponseModel | None = None
+    archived_agents: list[str] | None = None
+    archived_feeds: list[str] | None = None
+    deactivated_subscriptions: list[str] | None = None
+    skipped_entities: list[SkippedEntityModel] | None = None
+
+
+class DraftActivateResultModel(K2BaseModel):
+    success: bool | None = None
+    created_agent_ids: list[str] | None = None
+    created_feed_ids: list[str] | None = None
+    created_corpus_ids: list[str] | None = None
+    created_subscription_ids: list[str] | None = None
+    removed_binding_ids: list[str] | None = None
+    updated_entity_ids: list[str] | None = None
+    issues: list[ValidationIssueModel] | None = None
+
+
+class FieldDiffModel(K2BaseModel):
+    field: str | None = None
+    spec_value: str | None = None
+    actual_value: str | None = None
+
+
+class EntityDiffModel(K2BaseModel):
+    entity_type: str | None = None
+    entity_id: str | None = None
+    entity_name: str | None = None
+    drift_status: str | None = None
+    field_diffs: list[FieldDiffModel] | None = None
+
+
+class DriftSummaryModel(K2BaseModel):
+    total_entities: int | None = None
+    in_sync: int | None = None
+    drifted: int | None = None
+    missing: int | None = None
+    untracked: int | None = None
+
+
+class DriftReportModel(K2BaseModel):
+    pipeline_spec_id: str | None = None
+    status: str | None = None
+    entity_diffs: list[EntityDiffModel] | None = None
+    summary: DriftSummaryModel | None = None
+
+
+class RefreshChangesModel(K2BaseModel):
+    updated_entities: list[str] | None = None
+    removed_entities: list[str] | None = None
+    unchanged_entities: list[str] | None = None
+
+
+class RefreshResultModel(K2BaseModel):
+    draft_id: str | None = None
+    changes: RefreshChangesModel | None = None
+
+
+class GraphNodeModel(K2BaseModel):
+    id: str | None = None
+    type: str | None = None
+    label: str | None = None
+    metadata: dict[str, Any] | None = None
+    status: str | None = None
+    binding_role: str | None = None
+
+
+class GraphEdgeModel(K2BaseModel):
+    source: str | None = None
+    target: str | None = None
+    type: str | None = None
+    metadata: dict[str, Any] | None = None
+
+
+class GraphResponseModel(K2BaseModel):
+    pipeline_spec_id: str | None = None
+    status: str | None = None
+    nodes: list[GraphNodeModel] | None = None
+    edges: list[GraphEdgeModel] | None = None

sdk/models/projects.py

@@ -0,0 +1,19 @@
+"""Pydantic models for project types."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from sdk.models._base import K2BaseModel
+
+
+class ProjectResponseModel(K2BaseModel):
+    id: str
+    name: str
+    org_id: str
+    graph_rag_policy: dict[str, Any] | None
+
+
+class ProjectListResponseModel(K2BaseModel):
+    projects: list[ProjectResponseModel]
+    total: int

sdk/models/search.py

@@ -0,0 +1,149 @@
+"""Pydantic models for search types."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from sdk.models._base import K2BaseModel
+
+
+class SearchHybridConfigModel(K2BaseModel):
+    enabled: bool | None = None
+    fusion_mode: str | None = None
+    rrf_k: int | None = None
+    dense_weight: float | None = None
+    sparse_weight: float | None = None
+    metadata_sparse_enabled: bool | None = None
+
+
+class SearchRerankConfigModel(K2BaseModel):
+    enabled: bool | None = None
+    top_k: int | None = None
+
+
+class SearchGraphRagConfigModel(K2BaseModel):
+    enabled: bool | None = None
+    seed_chunks: int | None = None
+    seed_docs: int | None = None
+    max_entities_per_doc: int | None = None
+    candidate_scan_limit: int | None = None
+    min_entity_overlap: int | None = None
+    max_related_docs: int | None = None
+    max_chunks_per_doc: int | None = None
+    max_docs_per_entity: int | None = None
+    max_expanded_chunks: int | None = None
+    timeout_ms: int | None = None
+    depth_penalty: float | None = None
+    score_boost: float | None = None
+    shadow_mode: bool | None = None
+
+
+class SearchReturnConfigModel(K2BaseModel):
+    include_text: bool | None = None
+    include_scores: bool | None = None
+    include_provenance: bool | None = None
+
+
+class MetadataFilterModel(K2BaseModel):
+    key: str | None = None
+    op: str | None = None
+    value: Any | None = None
+
+
+class MetadataFiltersModel(K2BaseModel):
+    filters: list[MetadataFilterModel] | None = None
+    condition: str | None = None
+
+
+class SearchGenerationConfigModel(K2BaseModel):
+    model: str | None = None
+    thinking_budget: int | None = None
+    temperature: float | None = None
+    max_tokens: int | None = None
+    context_top_k: int | None = None
+
+
+class SearchOptionsModel(K2BaseModel):
+    query: str | None = None
+    top_k: int | None = None
+    filters: dict[str, Any] | MetadataFiltersModel | None = None
+    hybrid: SearchHybridConfigModel | None = None
+    graph_rag: SearchGraphRagConfigModel | None = None
+    rerank: SearchRerankConfigModel | None = None
+    return_config: SearchReturnConfigModel | None = None
+
+
+class SearchGenerateOptionsModel(K2BaseModel):
+    query: str | None = None
+    top_k: int | None = None
+    filters: dict[str, Any] | MetadataFiltersModel | None = None
+    hybrid: SearchHybridConfigModel | None = None
+    graph_rag: SearchGraphRagConfigModel | None = None
+    rerank: SearchRerankConfigModel | None = None
+    return_config: SearchReturnConfigModel | None = None
+    generation: SearchGenerationConfigModel | None = None
+
+
+class SearchBatchOptionsModel(K2BaseModel):
+    queries: list[str] | None = None
+    top_k: int | None = None
+    filters: dict[str, Any] | MetadataFiltersModel | None = None
+    hybrid: SearchHybridConfigModel | None = None
+    graph_rag: SearchGraphRagConfigModel | None = None
+    rerank: SearchRerankConfigModel | None = None
+    return_config: SearchReturnConfigModel | None = None
+
+
+class SearchResultModel(K2BaseModel):
+    chunk_id: str | None = None
+    score: float | None = None
+    raw_score: float | None = None
+    text: str | None = None
+    custom_metadata: dict | None = None
+    system_metadata: dict | None = None
+    offset_start: int | None = None
+    offset_end: int | None = None
+    page_start: int | None = None
+    page_end: int | None = None
+
+
+class SearchColdStartMetaModel(K2BaseModel):
+    likely: bool | None = None
+    dense_cache_hit: bool | None = None
+    sparse_cache_hit: bool | None = None
+    graph_cache_hit: bool | None = None
+
+
+class SearchGraphRagMetaModel(K2BaseModel):
+    enabled: bool | None = None
+    applied: bool | None = None
+    shadow_mode: bool | None = None
+    expanded_chunks: int | None = None
+    related_docs: int | None = None
+    source: str | None = None
+    fallback_used: bool | None = None
+    fallback_reason: str | None = None
+    timed_out: bool | None = None
+
+
+class SearchMetaModel(K2BaseModel):
+    cold_start: SearchColdStartMetaModel | None = None
+    graph_rag: SearchGraphRagMetaModel | None = None
+
+
+class SearchResponseModel(K2BaseModel):
+    results: list[SearchResultModel] | None = None
+    meta: SearchMetaModel | None = None
+
+
+class SearchBatchResponseModel(K2BaseModel):
+    responses: list[SearchResponseModel] | None = None
+
+
+class SearchGenerateResponseModel(K2BaseModel):
+    answer: str | None = None
+    model: str | None = None
+    thinking_budget: int | None = None
+    results: list[SearchResultModel] | None = None
+    meta: SearchMetaModel | None = None
+    used_sources: list[str] | None = None

sdk/models/training.py

@@ -0,0 +1,57 @@
+"""Pydantic models for training types."""
+
+from __future__ import annotations
+
+from sdk.models._base import K2BaseModel
+
+
+class TrainingDataBuildResponseModel(K2BaseModel):
+    job_id: str
+    dataset_id: str
+
+
+class TrainingDatasetResponseModel(K2BaseModel):
+    id: str | None = None
+    corpus_id: str | None = None
+    sample_count: int | None = None
+    dataset_hash: str | None = None
+    created_at: str | None = None
+
+
+class TrainingDatasetListResponseModel(K2BaseModel):
+    datasets: list[TrainingDatasetResponseModel]
+
+
+class TuningRunResponseModel(K2BaseModel):
+    run_id: str
+    status: str
+    job_id: str
+
+
+class TuningRunBuildResponseModel(K2BaseModel):
+    run_id: str
+    status: str
+    build_job_id: str
+    dataset_id: str
+
+
+class TuningRunDetailResponseModel(K2BaseModel):
+    run_id: str | None = None
+    status: str | None = None
+    metrics: dict | None = None
+
+
+class TuningRunLogsResponseModel(K2BaseModel):
+    lines: list[str]
+
+
+class TuningRunListResponseModel(K2BaseModel):
+    runs: list[TuningRunDetailResponseModel]
+
+
+class PromoteResponseModel(K2BaseModel):
+    model_id: str
+
+
+class CancelTuningRunResponseModel(K2BaseModel):
+    status: str