diff --git a/README.md b/README.md index c15d1c9..3d8ee2a 100644 --- a/README.md +++ b/README.md @@ -78,8 +78,8 @@ Make some changes to your git repo and stage them. Then, run `mf commit`! You sh 1 file changed, 14 insertions(+) ``` -### Create PRs With GPT Titles And Body -Make some changes to your branch and stage, and then commit them. Then, run `mf pr`! A PR should be created with a title and body generated by GPT, and a link to the PR should be printed to the console. +### Create PRs/MRs With GPT Titles And Body +Make some changes to your branch and stage, and then commit them. Then, run `mf pr` for GitHub or `mf mr` for GitLab! A pull request/merge request should be created with a title and body generated by GPT, and a link to the PR should be printed to the console. - To use this feature, you must first install and authenticate the [GitHub CLI](https://cli.github.com/). ## How does it work? diff --git a/mindflow/__init__.py b/mindflow/__init__.py index df0ed33..8a3be2e 100644 --- a/mindflow/__init__.py +++ b/mindflow/__init__.py @@ -1 +1 @@ -__version__ = "0.3.12" +__version__ = "0.3.13" diff --git a/mindflow/core/chat.py b/mindflow/core/chat.py index c37dcee..dbdcde7 100644 --- a/mindflow/core/chat.py +++ b/mindflow/core/chat.py @@ -1,4 +1,8 @@ +from typing import Optional from mindflow.settings import Settings +from mindflow.utils.constants import MinimumReservedLength +from mindflow.utils.prompts import CHAT_PROMPT_PREFIX +from mindflow.utils.token import get_token_count def run_chat(prompt: str) -> str: @@ -6,14 +10,26 @@ def run_chat(prompt: str) -> str: This function is used to generate a prompt and then use it as a prompt for GPT bot. """ settings = Settings() + completion_model = settings.mindflow_models.query.model + + if ( + get_token_count(completion_model, CHAT_PROMPT_PREFIX + prompt) + > completion_model.hard_token_limit - MinimumReservedLength.CHAT.value + ): + print("The prompt is too long. Please try again with a shorter prompt.") + return "" + # Prompt GPT through Mindflow API or locally - response: str = settings.mindflow_models.query.model( + response: Optional[str] = completion_model( [ { "role": "system", - "content": "You are a helpful virtual assistant responding to a users query using your general knowledge and the text provided below.", + "content": CHAT_PROMPT_PREFIX, }, {"role": "user", "content": prompt}, ] ) + + if response is None: + return "Unable to generate response. Please try again. If the problem persists, please raise an issue at: https://github.com/nollied/mindflow-cli/issues." return response diff --git a/mindflow/core/git/commit.py b/mindflow/core/git/commit.py index d3c0be9..ff64368 100644 --- a/mindflow/core/git/commit.py +++ b/mindflow/core/git/commit.py @@ -1,8 +1,9 @@ import subprocess -from typing import Tuple, Optional +from typing import Tuple, Optional, Union from mindflow.core.git.diff import run_diff from mindflow.settings import Settings +from mindflow.utils.errors import ModelError from mindflow.utils.prompt_builders import build_context_prompt from mindflow.utils.prompts import COMMIT_PROMPT_PREFIX @@ -20,9 +21,11 @@ def run_commit(args: Tuple[str], message_overwrite: Optional[str] = None) -> str if diff_output == "No staged changes.": return diff_output - response: str = settings.mindflow_models.query.model( + response: Union[ModelError, str] = settings.mindflow_models.query.model( build_context_prompt(COMMIT_PROMPT_PREFIX, diff_output) ) + if isinstance(response, ModelError): + return response.commit_message # add co-authorship to commit message response += "\n\nCo-authored-by: MindFlow " diff --git a/mindflow/core/git/diff.py b/mindflow/core/git/diff.py index eddaa41..4ddca4e 100644 --- a/mindflow/core/git/diff.py +++ b/mindflow/core/git/diff.py @@ -3,17 +3,19 @@ """ import concurrent.futures import subprocess -from typing import Dict +from typing import Dict, Union from typing import List -from typing import Optional from typing import Tuple from mindflow.db.objects.model import ConfiguredModel from mindflow.settings import Settings +from mindflow.utils.constants import MinimumReservedLength +from mindflow.utils.errors import ModelError from mindflow.utils.prompt_builders import build_context_prompt from mindflow.utils.prompts import GIT_DIFF_PROMPT_PREFIX -from mindflow.utils.diff_parser import parse_git_diff, IGNORE_FILE_EXTENSIONS +from mindflow.utils.diff_parser import parse_git_diff +from mindflow.utils.token import get_token_count def run_diff(args: Tuple[str]) -> str: @@ -35,18 +37,19 @@ def run_diff(args: Tuple[str]) -> str: if len(diff_dict) <= 0: return "No staged changes." - batched_parsed_diff_result = batch_git_diffs( - diff_dict, token_limit=completion_model.hard_token_limit - ) + batched_parsed_diff_result = batch_git_diffs(diff_dict, completion_model) - response: str = "" + diff_summary: str = "" if len(batched_parsed_diff_result) == 1: content = "" for file_name, diff_content in batched_parsed_diff_result[0]: content += f"*{file_name}*\n DIFF CONTENT: {diff_content}\n\n" - response = completion_model( + diff_response: Union[ModelError, str] = completion_model( build_context_prompt(GIT_DIFF_PROMPT_PREFIX, content) ) + if isinstance(diff_response, ModelError): + return diff_response.diff_message + diff_summary += diff_response else: with concurrent.futures.ThreadPoolExecutor() as executor: futures = [] @@ -62,43 +65,73 @@ def run_diff(args: Tuple[str]) -> str: # Process the results as they become available for future in concurrent.futures.as_completed(futures): - response += future.result() + diff_partial_response: Union[ModelError, str] = future.result() + if isinstance(diff_partial_response, ModelError): + return diff_partial_response.diff_partial_message + + diff_summary += diff_partial_response if len(excluded_filenames) > 0: - response += f"\n\nNOTE: The following files were excluded from the diff: {', '.join(excluded_filenames)}" + diff_summary += f"\n\nNOTE: The following files were excluded from the diff: {', '.join(excluded_filenames)}" - return response + return diff_summary import re def batch_git_diffs( - file_diffs: Dict[str, str], token_limit: int + file_diffs: Dict[str, str], model: ConfiguredModel ) -> List[List[Tuple[str, str]]]: batches = [] current_batch: List = [] - current_batch_size = 0 + current_batch_text = "" for file_name, diff_content in file_diffs.items(): - if len(diff_content) > token_limit: - chunks = [ - diff_content[i : i + token_limit] - for i in range(0, len(diff_content), token_limit) - ] + if ( + get_token_count(model, diff_content) + > model.hard_token_limit - MinimumReservedLength.DIFF.value + ): + ## Split the diff into chunks that are less than the token limit + chunks = [diff_content] + while True: + new_chunks = [] + for chunk in chunks: + if ( + get_token_count(model, chunk) + > model.hard_token_limit - MinimumReservedLength.DIFF.value + ): + half_len = len(chunk) // 2 + left_half = chunk[:half_len] + right_half = chunk[half_len:] + new_chunks.extend([left_half, right_half]) + else: + new_chunks.append(chunk) + if new_chunks == chunks: + break + chunks = new_chunks + + ## Add the chunks to the batch or multiple batches for chunk in chunks: - if current_batch_size + len(chunk) > token_limit * 2: + if ( + get_token_count(model, current_batch_text + chunk) + > model.hard_token_limit - MinimumReservedLength.DIFF.value + ): batches.append(current_batch) current_batch = [] - current_batch_size = 0 + current_batch_text = "" current_batch.append((file_name, chunk)) - current_batch_size += len(chunk) - elif current_batch_size + len(diff_content) > token_limit * 2: + current_batch_text += chunk + + elif ( + get_token_count(model, current_batch_text + diff_content) + > model.hard_token_limit - MinimumReservedLength.DIFF.value + ): batches.append(current_batch) current_batch = [(file_name, diff_content)] - current_batch_size = len(diff_content) + current_batch_text = diff_content else: current_batch.append((file_name, diff_content)) - current_batch_size += len(diff_content) + current_batch_text += diff_content if current_batch: batches.append(current_batch) return batches diff --git a/mindflow/core/git/mr.py b/mindflow/core/git/mr.py index 96bf8f7..bead0cd 100644 --- a/mindflow/core/git/mr.py +++ b/mindflow/core/git/mr.py @@ -32,7 +32,12 @@ def run_mr( return if not title or not description: - title, description = create_title_and_body(base_branch, title, description) + tital_description_tuple = create_title_and_body(base_branch, title, description) + + if not tital_description_tuple: + return + + title, description = tital_description_tuple create_merge_request(args, title, description) diff --git a/mindflow/core/git/pr.py b/mindflow/core/git/pr.py index 5b6df5f..0d7ff41 100644 --- a/mindflow/core/git/pr.py +++ b/mindflow/core/git/pr.py @@ -1,10 +1,11 @@ import concurrent.futures import subprocess -from typing import List, Optional, Tuple +from typing import List, Optional, Tuple, Union from mindflow.core.git.diff import run_diff from mindflow.settings import Settings from mindflow.utils.command_parse import get_flag_value +from mindflow.utils.errors import ModelError from mindflow.utils.prompt_builders import build_context_prompt from mindflow.utils.prompts import PR_BODY_PREFIX from mindflow.utils.prompts import PR_TITLE_PREFIX @@ -35,8 +36,12 @@ def run_pr(args: Tuple[str], title: Optional[str] = None, body: Optional[str] = return if not title or not body: - title, body = create_title_and_body(base_branch, title, body) + tital_body_tuple = create_title_and_body(base_branch, title, body) + if not tital_body_tuple: + return + + title, body = tital_body_tuple create_pull_request(args, title, body) @@ -58,11 +63,13 @@ def is_valid_pr(head_branch: str, base_branch: str) -> bool: def create_title_and_body( base_branch, title: Optional[str], body: Optional[str] -) -> Tuple[str, str]: +) -> Optional[Tuple[str, str]]: settings = Settings() diff_output = run_diff((base_branch,)) + title_response: Union[ModelError, str] + body_response: Union[ModelError, str] if title is None and body is None: pr_title_prompt = build_context_prompt(PR_TITLE_PREFIX, diff_output) pr_body_prompt = build_context_prompt(PR_BODY_PREFIX, diff_output) @@ -75,16 +82,25 @@ def create_title_and_body( settings.mindflow_models.query.model, pr_body_prompt ) - title = future_title.result() - body = future_body.result() + title_response = future_title.result() + body_response = future_body.result() else: if title is None: pr_title_prompt = build_context_prompt(PR_TITLE_PREFIX, diff_output) - title = settings.mindflow_models.query.model(pr_title_prompt) + title_response = settings.mindflow_models.query.model(pr_title_prompt) if body is None: pr_body_prompt = build_context_prompt(PR_BODY_PREFIX, diff_output) - body = settings.mindflow_models.query.model(pr_body_prompt) + body_response = settings.mindflow_models.query.model(pr_body_prompt) + + if isinstance(title_response, ModelError): + print(title_response.pr_message) + return None + if isinstance(body_response, ModelError): + print(body_response.pr_message) + return None + title = title if title is not None else title_response + body = body if body is not None else body_response return title, body diff --git a/mindflow/core/index.py b/mindflow/core/index.py index d354ab2..2eacb44 100644 --- a/mindflow/core/index.py +++ b/mindflow/core/index.py @@ -1,10 +1,10 @@ """ `generate` command """ -from asyncio import Future from concurrent.futures import ThreadPoolExecutor from copy import deepcopy -from typing import List +import logging +from typing import List, Union from typing import Optional import numpy as np @@ -18,8 +18,10 @@ from mindflow.resolving.resolve import resolve_all from mindflow.resolving.resolve import return_if_indexable from mindflow.settings import Settings +from mindflow.utils.errors import ModelError from mindflow.utils.prompt_builders import build_context_prompt from mindflow.utils.prompts import INDEX_PROMPT_PREFIX +from mindflow.utils.token import get_batch_token_count, get_token_count def run_index(document_paths: List[str], refresh: bool, force: bool) -> None: @@ -97,9 +99,14 @@ def __init__( self.start = start self.end = end if text: - self.summary = completion_model( + response: Union[str, ModelError] = completion_model( build_context_prompt(INDEX_PROMPT_PREFIX, text) ) + if isinstance(response, ModelError): + self.summary = "" + print(response.index_message) + else: + self.summary = response def set_leaves(self, leaves: List["Node"]) -> None: self.leaves = leaves @@ -136,15 +143,6 @@ def iterative_to_dict(self) -> dict: return node_dict -def count_tokens(text: str) -> int: - """ - Counts/estimates the number of tokens this text will consume by GPT. - """ - # tokenizer = GPT2TokenizerFast.from_pretrained("gpt2") - # count = len(tokenizer(text)['input_ids']) - return len(text) // 4 # Token Estimation for speed - - # This function is used to split a string into chunks of a specified token limit using binary search def binary_split_raw_text_to_nodes( completion_model: ConfiguredModel, text: str @@ -156,7 +154,10 @@ def binary_split_raw_text_to_nodes( stack = [(0, len(text))] while stack: start, end = stack.pop() - if count_tokens(text[start:end]) < completion_model.soft_token_limit: + if ( + get_token_count(completion_model, text[start:end]) + < completion_model.soft_token_limit + ): nodes.append(Node(completion_model, start, end, text[start:end])) else: mid = ((end - start) // 2) + start @@ -176,7 +177,9 @@ def binary_split_nodes_to_chunks( while stack: nodes, start, end = stack.pop() if ( - sum(count_tokens(node.summary) for node in nodes[start:end]) + get_batch_token_count( + completion_model, [node.summary for node in nodes[start:end]] + ) < completion_model.soft_token_limit ): chunks.append(nodes[start:end]) @@ -195,7 +198,10 @@ def create_nodes(completion_model: ConfiguredModel, leaf_nodes: List[Node]) -> N while stack: leaf_nodes, start, end = stack.pop() if ( - sum(count_tokens(leaf_node.summary) for leaf_node in leaf_nodes[start:end]) + get_batch_token_count( + completion_model, + [leaf_node.summary for leaf_node in leaf_nodes[start:end]], + ) > completion_model.soft_token_limit ): node_chunks: List[List[Node]] = binary_split_nodes_to_chunks( @@ -222,7 +228,7 @@ def create_text_search_tree(completion_model: ConfiguredModel, text: str) -> dic """ This function is used to create a tree of responses from the OpenAI API """ - if count_tokens(text) < completion_model.soft_token_limit: + if get_token_count(completion_model, text) < completion_model.soft_token_limit: return Node(completion_model, 0, len(text), text).to_dict() return create_nodes( diff --git a/mindflow/core/query.py b/mindflow/core/query.py index d5075a4..56ec2b6 100644 --- a/mindflow/core/query.py +++ b/mindflow/core/query.py @@ -4,7 +4,7 @@ import sys from concurrent.futures import as_completed from concurrent.futures import ThreadPoolExecutor -from typing import Dict +from typing import Dict, Union from typing import List from typing import Optional from typing import Tuple @@ -17,6 +17,9 @@ from mindflow.db.objects.model import ConfiguredModel from mindflow.resolving.resolve import resolve_all from mindflow.settings import Settings +from mindflow.utils.constants import MinimumReservedLength +from mindflow.utils.errors import ModelError +from mindflow.utils.token import get_token_count def run_query(document_paths: List[str], query: str): @@ -33,11 +36,14 @@ def run_query(document_paths: List[str], query: str): select_content( query, document_references, - completion_model.hard_token_limit, + completion_model, embedding_model, ), ) - response = completion_model(messages) + response: Union[ModelError, str] = completion_model(messages) + if isinstance(response, ModelError): + return response.query_message + return response @@ -58,7 +64,7 @@ def build_query_messages(query: str, content: str) -> List[Dict]: def select_content( query: str, document_references: List[DocumentReference], - token_limit: int, + completion_model: ConfiguredModel, embedding_model: ConfiguredModel, ) -> str: """ @@ -73,7 +79,9 @@ def select_content( ) sys.exit(1) - selected_content = trim_content(embedding_ranked_document_chunks, token_limit) + selected_content = trim_content( + embedding_ranked_document_chunks, completion_model, query + ) return selected_content @@ -109,9 +117,14 @@ def from_search_tree( document.search_tree["end"], ) ] - embeddings: List[np.ndarray] = [ - embedding_model(document.search_tree["summary"]) - ] + embedding_response: Union[ModelError, np.ndarray] = embedding_model( + document.search_tree["summary"] + ) + if isinstance(embedding_response, ModelError): + print(embedding_response.embedding_message) + return [], [] + + embeddings: List[np.ndarray] = [embedding_response] rolling_summary: List[str] = [] while stack: node = stack.pop() @@ -120,32 +133,48 @@ def from_search_tree( for leaf in node["leaves"]: stack.append(leaf) chunks.append(cls(document.path, leaf["start"], leaf["end"])) - rolling_summary_embedding = embedding_model( + rolling_summary_embedding_response: Union[ + np.ndarray, ModelError + ] = embedding_model( "\n\n".join(rolling_summary) + "\n\n" + leaf["summary"], ) - embeddings.append(rolling_summary_embedding) + if isinstance(rolling_summary_embedding_response, ModelError): + print(rolling_summary_embedding_response.embedding_message) + continue + embeddings.append(rolling_summary_embedding_response) rolling_summary.pop() return chunks, embeddings -def trim_content(ranked_document_chunks: List[DocumentChunk], token_limit: int) -> str: +def trim_content( + ranked_document_chunks: List[DocumentChunk], model: ConfiguredModel, query: str +) -> str: """ This function is used to select the most relevant content for the prompt. """ selected_content: str = "" - char_limit: int = token_limit * 3 for document_chunk in ranked_document_chunks: if document_chunk: with open(document_chunk.path, "r", encoding="utf-8") as file: file.seek(document_chunk.start) text = file.read(document_chunk.end - document_chunk.start) - if len(selected_content + text) > char_limit: - selected_content += text[: char_limit - len(selected_content)] - break - selected_content += text + # Perform a binary search to find the maximum amount of text that fits within the token limit + left, right = 0, len(text) + while left <= right: + mid = (left + right) // 2 + if ( + get_token_count(model, query + selected_content + text[:mid]) + <= model.hard_token_limit - MinimumReservedLength.QUERY.value + ): + left = mid + 1 + else: + right = mid - 1 + + # Add the selected text to the selected content + selected_content += text[:right] return selected_content diff --git a/mindflow/db/db/static.py b/mindflow/db/db/static.py index c32fae3..d1d409a 100644 --- a/mindflow/db/db/static.py +++ b/mindflow/db/db/static.py @@ -1,5 +1,3 @@ -from typing import Union - from mindflow.db.db.database import Collection from mindflow.db.db.database import Database from mindflow.db.objects.static_definition.mind_flow_model import MINDFLOW_MODEL_STATIC diff --git a/mindflow/db/objects/model.py b/mindflow/db/objects/model.py index ecc7827..926c6c9 100644 --- a/mindflow/db/objects/model.py +++ b/mindflow/db/objects/model.py @@ -1,14 +1,26 @@ -from typing import Optional +from typing import Optional, Union import openai +import numpy as np from traitlets import Callable +try: + import tiktoken +except ImportError: + print( + "tiktoken not not available in python<=v3.8. Estimation of tokens will be less precise, which may impact performance and quality of responses." + ) + print("Upgrade to python v3.8 or higher for better results.") + pass + + from mindflow.db.db.database import Collection from mindflow.db.objects.base import BaseObject from mindflow.db.objects.base import StaticObject from mindflow.db.objects.service import ServiceConfig from mindflow.db.objects.static_definition.model import ModelID from mindflow.db.objects.static_definition.service import ServiceID +from mindflow.utils.errors import ModelError class Model(StaticObject): @@ -43,6 +55,12 @@ class ConfiguredModel(Callable): name: str service: str model_type: str + + try: + tokenizer: tiktoken.Encoding + except NameError: + pass + hard_token_limit: int token_cost: int token_cost_unit: str @@ -64,30 +82,42 @@ def __init__(self, model_id: str): if value not in [None, ""]: setattr(self, key, value) + try: + if self.service == ServiceID.OPENAI.value: + self.tokenizer = tiktoken.encoding_for_model(self.id) + except NameError: + pass + service_config = ServiceConfig.load(f"{self.service}_config") self.api_key = service_config.api_key def openai_chat_completion( self, messages: list, - max_tokens: int = 500, temperature: float = 0.0, + max_tokens: Optional[int] = None, stop: Optional[list] = None, - ): - openai.api_key = self.api_key - return openai.ChatCompletion.create( - model=self.id, - messages=messages, - temperature=temperature, - max_tokens=max_tokens, - stop=stop, - )["choices"][0]["message"]["content"] - - def openai_embedding(self, text: str): - openai.api_key = self.api_key - return openai.Embedding.create(engine=self.id, input=text)["data"][0][ - "embedding" - ] + ) -> Union[str, ModelError]: + try: + openai.api_key = self.api_key + return openai.ChatCompletion.create( + model=self.id, + messages=messages, + temperature=temperature, + max_tokens=max_tokens, + stop=stop, + )["choices"][0]["message"]["content"] + except ModelError as e: + return e + + def openai_embedding(self, text: str) -> Union[np.ndarray, ModelError]: + try: + openai.api_key = self.api_key + return openai.Embedding.create(engine=self.id, input=text)["data"][0][ + "embedding" + ] + except ModelError as e: + return e def __call__(self, prompt, *args, **kwargs): if self.service == ServiceID.OPENAI.value: diff --git a/mindflow/db/objects/static_definition/mind_flow_model.py b/mindflow/db/objects/static_definition/mind_flow_model.py index 94b904c..a4bfd1e 100644 --- a/mindflow/db/objects/static_definition/mind_flow_model.py +++ b/mindflow/db/objects/static_definition/mind_flow_model.py @@ -1,6 +1,3 @@ -from typing import Type -from typing import Union - from mindflow.db.objects.static_definition.model import ModelID from mindflow.db.objects.static_definition.model_type import ModelType from mindflow.utils.enum import ExtendedEnum diff --git a/mindflow/db/objects/static_definition/model.py b/mindflow/db/objects/static_definition/model.py index e261356..fb63d6c 100644 --- a/mindflow/db/objects/static_definition/model.py +++ b/mindflow/db/objects/static_definition/model.py @@ -1,6 +1,3 @@ -from typing import Type -from typing import Union - from mindflow.db.objects.static_definition.model_type import ModelType from mindflow.utils.enum import ExtendedEnum diff --git a/mindflow/db/objects/static_definition/model_type.py b/mindflow/db/objects/static_definition/model_type.py index b0a4d1c..7e74c73 100644 --- a/mindflow/db/objects/static_definition/model_type.py +++ b/mindflow/db/objects/static_definition/model_type.py @@ -1,5 +1,3 @@ -from typing import Union - from mindflow.utils.enum import ExtendedEnum diff --git a/mindflow/db/objects/static_definition/object.py b/mindflow/db/objects/static_definition/object.py index 21cbd25..70d8492 100644 --- a/mindflow/db/objects/static_definition/object.py +++ b/mindflow/db/objects/static_definition/object.py @@ -1,6 +1,3 @@ -from typing import Type -from typing import Union - from mindflow.db.objects.document import Document from mindflow.db.objects.mindflow_model import MindFlowModel from mindflow.db.objects.mindflow_model import MindFlowModelConfig diff --git a/mindflow/db/objects/static_definition/service.py b/mindflow/db/objects/static_definition/service.py index abbb3ba..cb57d31 100644 --- a/mindflow/db/objects/static_definition/service.py +++ b/mindflow/db/objects/static_definition/service.py @@ -1,6 +1,3 @@ -from typing import Type -from typing import Union - from mindflow.db.objects.static_definition.model import ModelID from mindflow.db.objects.static_definition.model import ModelOpenAI from mindflow.db.objects.static_definition.model import ModelTextCompletionOpenAI diff --git a/mindflow/resolving/resolvers/base_resolver.py b/mindflow/resolving/resolvers/base_resolver.py index 371557f..0de7f0b 100644 --- a/mindflow/resolving/resolvers/base_resolver.py +++ b/mindflow/resolving/resolvers/base_resolver.py @@ -2,7 +2,6 @@ Base Resolver Class """ from typing import List -from typing import Optional from mindflow.db.objects.document import DocumentReference @@ -12,12 +11,6 @@ class BaseResolver: Base class for resolvers """ - @staticmethod - def read_document(document_path: str) -> Optional[str]: - """ - Read a document. - """ - @staticmethod def should_resolve(document_path: str) -> bool: """ diff --git a/mindflow/test.py b/mindflow/test.py new file mode 100644 index 0000000..83f5199 --- /dev/null +++ b/mindflow/test.py @@ -0,0 +1,5 @@ +import tiktoken + +tokenizer = tiktoken.encoding_for_model("gpt-3.5-turbo") + +print(tokenizer.encode_batch(["Hello, world!", "Whats up?"])) diff --git a/mindflow/utils/command_parse.py b/mindflow/utils/command_parse.py index 710842d..e72b34c 100644 --- a/mindflow/utils/command_parse.py +++ b/mindflow/utils/command_parse.py @@ -15,6 +15,9 @@ def get_flag_value(args: Tuple[str], flag: List[str]) -> Optional[str]: def get_flag_bool(args: Tuple[str], flag: str) -> bool: + """ + Returns True if the flag is in the list of arguments. + """ try: return args.index(flag) >= 0 except: diff --git a/mindflow/utils/constants.py b/mindflow/utils/constants.py new file mode 100644 index 0000000..21791b8 --- /dev/null +++ b/mindflow/utils/constants.py @@ -0,0 +1,7 @@ +from enum import Enum + + +class MinimumReservedLength(Enum): + CHAT = 500 + QUERY = 500 + DIFF = 600 diff --git a/mindflow/utils/errors.py b/mindflow/utils/errors.py new file mode 100644 index 0000000..5918d2f --- /dev/null +++ b/mindflow/utils/errors.py @@ -0,0 +1,50 @@ +GITHUB_ISSUE_MESSAGE = "If the problem persists, please raise an issue at: https://github.com/nollied/mindflow-cli/issues" +CONNECTION_MESSAGE = "Please check your internet connection and try again." + + +class ModelError(Exception): + """Base class for all exceptions raised by this module.""" + + def __init__(self, message): + self.message = message + super().__init__(self.message) + + @property + def base_message(self): + return f"Model API failed to return response for chat/query. {CONNECTION_MESSAGE}. {GITHUB_ISSUE_MESSAGE}." + + @property + def commit_message(self): + return f"Model API failed to return response for commit. {CONNECTION_MESSAGE}. {GITHUB_ISSUE_MESSAGE}" + + @property + def diff_message(self): + return f"Model API failed to return response for diff. {CONNECTION_MESSAGE}. {GITHUB_ISSUE_MESSAGE}" + + @property + def diff_partial_message(self): + return f"Warning: model API failed to return response for part of, or entire, diff. {CONNECTION_MESSAGE}. {GITHUB_ISSUE_MESSAGE}" + + @property + def pr_message(self): + return f"Model API failed to return response for pr/mr. {CONNECTION_MESSAGE}. {GITHUB_ISSUE_MESSAGE}" + + @property + def index_message(self): + return f"Warning: Model API failed to return response for a document chunk. {CONNECTION_MESSAGE}. {GITHUB_ISSUE_MESSAGE}" + + @property + def query_message(self): + return f"Model API failed to return response for query. {CONNECTION_MESSAGE}. {GITHUB_ISSUE_MESSAGE}" + + @property + def embedding_message(self): + return f"Warning: Model API failed to return response for embedding. {CONNECTION_MESSAGE}. {GITHUB_ISSUE_MESSAGE}" + + +class EmbeddingModelError(Exception): + """Base class for all exceptions raised by this module.""" + + def __init__(self, message): + self.message = message + super().__init__(self.message) diff --git a/mindflow/utils/prompts.py b/mindflow/utils/prompts.py index cc1916e..8b1be87 100644 --- a/mindflow/utils/prompts.py +++ b/mindflow/utils/prompts.py @@ -16,6 +16,7 @@ yet condensed string that can serve as an index for the contents or the purpose \ of a file. I want you to respond in as few words as possible while still conveying \ the full content and purpose of this file." +CHAT_PROMPT_PREFIX = "You are a helpful virtual assistant responding to a users query using your general knowledge and the text provided below." COMMIT_PROMPT_PREFIX = "Please provide a commit message for the following changes. Only respond with the commit message and nothing else." PR_TITLE_PREFIX = "Please provide a title for the following pull request using this git diff summary. Only respond with the title and nothing else." PR_BODY_PREFIX = "Please provide a body for the following pull request using this git diff summary. I want you to keep it high level, and give core \ diff --git a/mindflow/utils/token.py b/mindflow/utils/token.py new file mode 100644 index 0000000..a7e27d1 --- /dev/null +++ b/mindflow/utils/token.py @@ -0,0 +1,22 @@ +from typing import List +from mindflow.db.objects.model import ConfiguredModel + + +def get_token_count(model: ConfiguredModel, text: str) -> int: + """ + This function is used to get the token count of a string. + """ + try: + return len(model.tokenizer.encode(text)) + except Exception: + return len(text) // 3 + + +def get_batch_token_count(model: ConfiguredModel, texts: List[str]) -> int: + """ + This function is used to get the token count of a list of strings. + """ + try: + return sum([len(encoding) for encoding in model.tokenizer.encode_batch(texts)]) + except Exception: + return sum([len(text) // 3 for text in texts]) diff --git a/requirements.txt b/requirements.txt index f7c5ba5..16befaa 100644 --- a/requirements.txt +++ b/requirements.txt @@ -2,6 +2,7 @@ alive-progress click numpy openai==0.27.0 +tiktoken==0.3.0; python_version >= '3.8' pytest scikit-learn tqdm