Skip to content

dependencies

CohereChatModel

Bases: ChatLLMModel

Cohere chat model.

Source code in api/dependencies.py 
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
class CohereChatModel(ChatLLMModel):
    """Cohere chat model."""

    def chat(self, chat_input) -> str:
        """
        Chat with the model.

        Args:
            chat_input (core.models.chat.ChatInput): Chat input

        Returns:
            str: Chat response
        """
        messages = [transform_chat_message(m) for m in chat_input.messages]
        res = co.chat(messages=messages, model="command-r-plus-08-2024")
        return res.message.content[0].text

chat(chat_input)

Chat with the model.

Parameters:

Name Type Description Default
chat_input ChatInput

Chat input

 required

Returns:

Name Type Description
str str

Chat response
Source code in api/dependencies.py 
54
55
56
57
58
59
60
61
62
63
64
65
66
def chat(self, chat_input) -> str:
    """
    Chat with the model.

    Args:
        chat_input (core.models.chat.ChatInput): Chat input

    Returns:
        str: Chat response
    """
    messages = [transform_chat_message(m) for m in chat_input.messages]
    res = co.chat(messages=messages, model="command-r-plus-08-2024")
    return res.message.content[0].text

CohereEmbeddingsFunction

Bases: EmbeddingFunction

Cohere embeddings function.

Source code in api/dependencies.py 
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
class CohereEmbeddingsFunction(EmbeddingFunction):
    """Cohere embeddings function."""

    # skipcq: PYL-W0622
    def __call__(self, input: Documents) -> Optional[list[list[float]]]:
        """
        Call embeddings

        Args:
              input (Documents): embeddings input

        Returns:
            Optional[list[list[float]]: Embeddings output
        """
        response = co.embed(
            texts=input,
            model="embed-multilingual-v2.0",
            input_type="search_document",
            embedding_types=["float"],
        )
        return response.embeddings.float_

__call__(input)

Call embeddings

Parameters:

Name Type Description Default
input Documents

embeddings input

 required

Returns:

Type Description
Optional[list[list[float]]]

Optional[list[list[float]]: Embeddings output
Source code in api/dependencies.py 
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
def __call__(self, input: Documents) -> Optional[list[list[float]]]:
    """
    Call embeddings

    Args:
          input (Documents): embeddings input

    Returns:
        Optional[list[list[float]]: Embeddings output
    """
    response = co.embed(
        texts=input,
        model="embed-multilingual-v2.0",
        input_type="search_document",
        embedding_types=["float"],
    )
    return response.embeddings.float_

CohereRerankModel

Bases: RerankModel

Cohere rerank model.

Source code in api/dependencies.py 
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
class CohereRerankModel(RerankModel):
    """Cohere rerank model."""

    def rerank_documents(self, query, docs) -> list[float]:
        """
        Rerank the documents based on the query.

        Args:
            query (str): The query use to rerank
            docs (list[str]): List of documents

        Returns:
            List[float]: List of relevance scores
        """
        res = co.rerank(documents=docs, query=query, model="rerank-multilingual-v2.0")
        sorted_index = sorted(res.results, key=lambda x: x.index)
        return [el.relevance_score for el in sorted_index]

rerank_documents(query, docs)

Rerank the documents based on the query.

Parameters:

Name Type Description Default
query str

The query use to rerank

 required
docs list[str]

List of documents

 required

Returns:

Type Description
list[float]

List[float]: List of relevance scores
Source code in api/dependencies.py 
72
73
74
75
76
77
78
79
80
81
82
83
84
85
def rerank_documents(self, query, docs) -> list[float]:
    """
    Rerank the documents based on the query.

    Args:
        query (str): The query use to rerank
        docs (list[str]): List of documents

    Returns:
        List[float]: List of relevance scores
    """
    res = co.rerank(documents=docs, query=query, model="rerank-multilingual-v2.0")
    sorted_index = sorted(res.results, key=lambda x: x.index)
    return [el.relevance_score for el in sorted_index]

get_chat_model()

Creates a chat model.

Returns:

Name Type Description
ChatCohere CohereChatModel

Chat model
Source code in api/dependencies.py 
138
139
140
141
142
143
144
145
def get_chat_model() -> CohereChatModel:
    """
    Creates a chat model.

    Returns:
        ChatCohere: Chat model
    """
    return CohereChatModel()

get_chroma_client()

Creates a Chroma client to connect to the Chroma server.

Returns:

Type Description
Client

chromadb.Client: Chroma client
Source code in api/dependencies.py 
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
def get_chroma_client() -> chromadb.Client:
    """
    Creates a Chroma client to connect to the Chroma server.

    Returns:
        chromadb.Client: Chroma client
    """
    try:
        return chromadb.HttpClient(
            host=CHROMA_HOST,
            port=CHROMA_PORT,
        )
    except BaseException as e:
        logger.error("Error creating Chroma client: %s", e)
        return chromadb.Client()

get_embeddings_function()

Creates a Cohere embeddings function.

Returns:

Name Type Description
CohereEmbeddings CohereEmbeddingsFunction

Cohere embeddings function
Source code in api/dependencies.py 
105
106
107
108
109
110
111
112
113
114
def get_embeddings_function() -> CohereEmbeddingsFunction:
    """
    Creates a Cohere embeddings function.

    Returns:
        CohereEmbeddings: Cohere embeddings function
    """
    return CohereEmbeddingsFunction(
        cohere_api_key=COHERE_API_KEY,
    )

get_rerank_model()

Creates a rerank model.

Returns:

Name Type Description
CohereRerankModel CohereRerankModel

Rerank model
Source code in api/dependencies.py 
148
149
150
151
152
153
154
155
def get_rerank_model() -> CohereRerankModel:
    """
    Creates a rerank model.

    Returns:
        CohereRerankModel: Rerank model
    """
    return CohereRerankModel()

transform_chat_message(chat_message)

Transforms a chat message to a dictionary.

Parameters:

Name Type Description Default
chat_message ChatMessage

Chat message

 required

Returns:

Name Type Description
dict dict

Chat message dictionary
Source code in api/dependencies.py 
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
def transform_chat_message(chat_message: ChatMessage) -> dict:
    """
    Transforms a chat message to a dictionary.

    Args:
        chat_message (ChatMessage): Chat message

    Returns:
        dict: Chat message dictionary
    """
    role = "user"
    if chat_message.role == ChatMessageRole.Ai:
        role = "assistant"
    if chat_message.role == ChatMessageRole.System:
        role = "system"
    return {
        "role": role,
        "content": chat_message.content,
    }