pharia_skill/
csi.rs

1pub mod chunking;
2pub mod document_index;
3pub mod inference;
4pub mod language;
5
6use chunking::ChunkRequest;
7use document_index::{Document, SearchResult};
8use inference::{ChatRequest, ChatResponse, Completion, CompletionRequest};
9use language::SelectLanguageRequest;
10use serde::{Deserialize, Serialize};
11
12use crate::{DocumentPath, LanguageCode, SearchRequest};
13
14/// Cognitive System Interface
15pub trait Csi {
16    /// Chunk the given text into smaller pieces that fit within the
17    /// maximum token amount for a given model.
18    fn chunk(&self, request: ChunkRequest) -> Vec<String> {
19        self.chunk_concurrently(vec![request]).remove(0)
20    }
21
22    /// Process multiple chunking requests at once
23    fn chunk_concurrently(&self, requests: Vec<ChunkRequest>) -> Vec<Vec<String>>;
24
25    /// Search for documents in a given index.
26    fn search(&self, request: SearchRequest) -> Vec<SearchResult> {
27        self.search_concurrently(vec![request]).remove(0)
28    }
29
30    /// Process multiple search requests at once
31    fn search_concurrently(&self, requests: Vec<SearchRequest>) -> Vec<Vec<SearchResult>>;
32
33    /// Retrieve a document from the Document Index by its path.
34    ///
35    /// # Errors
36    /// Will return an error if document metadata cannot be deserialized.
37    fn document<Metadata>(&self, path: DocumentPath) -> anyhow::Result<Document<Metadata>>
38    where
39        Metadata: for<'a> Deserialize<'a> + Serialize,
40    {
41        Ok(self.documents(vec![path])?.remove(0))
42    }
43
44    /// Retrieve multiple documents from the Document Index by their paths.
45    ///
46    /// # Errors
47    /// Will return an error if document metadata cannot be deserialized.
48    fn documents<'m, Metadata>(
49        &self,
50        paths: Vec<DocumentPath>,
51    ) -> anyhow::Result<Vec<Document<Metadata>>>
52    where
53        Metadata: for<'a> Deserialize<'a> + Serialize;
54
55    /// Retrieve a document's metadata from the Document Index by its path.
56    ///
57    /// # Errors
58    /// Will return an error if metadata cannot be deserialized.
59    fn document_metadata<Metadata>(&self, path: DocumentPath) -> anyhow::Result<Option<Metadata>>
60    where
61        Metadata: for<'a> Deserialize<'a> + Serialize,
62    {
63        Ok(self.documents_metadata(vec![path])?.remove(0))
64    }
65
66    /// Retrieve multiple documents' metadata from the Document Index by their paths.
67    ///
68    /// # Errors
69    /// Will return an error if metadata cannot be deserialized.
70    fn documents_metadata<Metadata>(
71        &self,
72        paths: Vec<DocumentPath>,
73    ) -> anyhow::Result<Vec<Option<Metadata>>>
74    where
75        Metadata: for<'a> Deserialize<'a> + Serialize;
76
77    /// Send messages with a particular role to a model and receive a response.
78    /// Provides a higher level interface than completion for chat scenarios.
79    fn chat(&self, request: ChatRequest) -> ChatResponse {
80        self.chat_concurrently(vec![request]).remove(0)
81    }
82
83    /// Process multiple chat requests at once
84    fn chat_concurrently(&self, requests: Vec<ChatRequest>) -> Vec<ChatResponse>;
85
86    /// Generate a completion for a given prompt using a specific model.
87    fn complete(&self, request: CompletionRequest) -> Completion {
88        self.complete_concurrently(vec![request]).remove(0)
89    }
90
91    /// Process multiple completion requests at once
92    fn complete_concurrently(&self, requests: Vec<CompletionRequest>) -> Vec<Completion>;
93
94    /// Select the detected language for the provided input based on the list of possible languages.
95    /// If no language matches, None is returned.
96    ///
97    /// text: Text input
98    /// languages: All languages that should be considered during detection.
99    fn select_language(&self, request: SelectLanguageRequest) -> Option<LanguageCode> {
100        self.select_language_concurrently(vec![request]).remove(0)
101    }
102
103    /// Process multiple select language requests at once
104    fn select_language_concurrently(
105        &self,
106        requests: Vec<SelectLanguageRequest>,
107    ) -> Vec<Option<LanguageCode>>;
108}