fix: use the latest user messages block instead of single message

data/archived.jsonl

@@ -1,7 +1,7 @@
 {"name":"@prefix/archived-npm-dummy","type":"npm","description":"Dummy archived to test with encoded package name on npm"}
 {"name":"archived-npm-dummy","type":"npm","description":"Dummy archived to test with simple package name on npm"}
 {"name":"@prefix/archived-pypi-dummy","type":"pypi","description":"Dummy archived to test with encoded package name on pypi"}
-{"name":"archived-pypi-dummy","type":"pypi","description":"Dummy archived to test with simple package name on pypi"}
+{"name":"archived_pypi_dummy","type":"pypi","description":"Dummy archived to test with simple package name on pypi"}
 {"name":"@prefix/archived-maven-dummy","type":"maven","description":"Dummy archived to test with encoded package name on maven"}
 {"name":"archived-maven-dummy","type":"maven","description":"Dummy archived to test with simple package name on maven"}
 {"name":"github.com/archived-go-dummy","type":"npm","description":"Dummy archived to test with encoded package name on go"}

data/deprecated.jsonl

@@ -1,7 +1,7 @@
 {"name":"@prefix/deprecated-npm-dummy","type":"npm","description":"Dummy deprecated to test with encoded package name on npm"}
 {"name":"deprecated-npm-dummy","type":"npm","description":"Dummy deprecated to test with simple package name on npm"}
 {"name":"@prefix/deprecated-pypi-dummy","type":"pypi","description":"Dummy deprecated to test with encoded package name on pypi"}
-{"name":"deprecated-pypi-dummy","type":"pypi","description":"Dummy deprecated to test with simple package name on pypi"}
+{"name":"deprecated_pypi_dummy","type":"pypi","description":"Dummy deprecated to test with simple package name on pypi"}
 {"name":"@prefix/deprecated-maven-dummy","type":"maven","description":"Dummy deprecated to test with encoded package name on maven"}
 {"name":"deprecated-maven-dummy","type":"maven","description":"Dummy deprecated to test with simple package name on maven"}
 {"name":"github.com/deprecated-go-dummy","type":"npm","description":"Dummy deprecated to test with encoded package name on go"}

data/malicious.jsonl

@@ -1,7 +1,7 @@
 {"name":"@prefix/malicious-npm-dummy","type":"npm","description":"Dummy malicious to test with encoded package name on npm"}
 {"name":"malicious-npm-dummy","type":"npm","description":"Dummy malicious to test with simple package name on npm"}
 {"name":"@prefix/malicious-pypi-dummy","type":"pypi","description":"Dummy malicious to test with encoded package name on pypi"}
-{"name":"malicious-pypi-dummy","type":"pypi","description":"Dummy malicious to test with simple package name on pypi"}
+{"name":"malicious_pypi_dummy","type":"pypi","description":"Dummy malicious to test with simple package name on pypi"}
 {"name":"@prefix/malicious-maven-dummy","type":"maven","description":"Dummy malicious to test with encoded package name on maven"}
 {"name":"malicious-maven-dummy","type":"maven","description":"Dummy malicious to test with simple package name on maven"}
 {"name":"github.com/malicious-go-dummy","type":"go","description":"Dummy malicious to test with encoded package name on go"}

pyproject.toml

@@ -28,6 +28,7 @@ tree-sitter-python = ">=0.23.6"
 tree-sitter-rust = ">=0.23.2"
 sqlite-vec-sl-tmp = "^0.0.4"
 alembic = ">=1.14.0"
+pygments = "^2.19.1"
 
 [tool.poetry.group.dev.dependencies]
 pytest = ">=7.4.0"

src/codegate/pipeline/base.py

@@ -231,8 +231,54 @@ def get_last_user_message(
  return None
  for i in reversed(range(len(request["messages"]))):
  if request["messages"][i]["role"] == "user":
- content = request["messages"][i]["content"]
- return content, i
+ content = request["messages"][i]["content"] # type: ignore
+ return str(content), i
+
+ return None
+
+ @staticmethod
+ def get_last_user_message_block(
+ request: ChatCompletionRequest,
+ ) -> Optional[str]:
+ """
+ Get the last block of consecutive 'user' messages from the request.
+
+ Args:
+ request (ChatCompletionRequest): The chat completion request to process
+
+ Returns:
+ Optional[str]: A string containing all consecutive user messages in the
+ last user message block, separated by newlines, or None if
+ no user message block is found.
+ """
+ if request.get("messages") is None:
+ return None
+
+ user_messages = []
+ messages = request["messages"]
+
+ # Iterate in reverse to find the last block of consecutive 'user' messages
+ for i in reversed(range(len(messages))):
+ if messages[i]["role"] == "user" or messages[i]["role"] == "assistant":
+ content_str = None
+ if "content" in messages[i]:
+ content_str = messages[i]["content"] # type: ignore
+ else:
+ continue
+
+ if messages[i]["role"] == "user":
+ user_messages.append(content_str)
+ # specifically for Aider, when "ok." block is found, stop
+ if content_str == "Ok." and messages[i]["role"] == "assistant":
+ break
+ else:
+ # Stop when a message with a different role is encountered
+ if user_messages:
+ break
+
+ # Reverse the collected user messages to preserve the original order
+ if user_messages:
+ return "\n".join(reversed(user_messages))
 
  return None
 

src/codegate/pipeline/codegate_context_retriever/codegate.py

@@ -59,44 +59,53 @@ async def process(
  """
  Use RAG DB to add context to the user request
  """
- # Get the latest user messages
- user_messages = self.get_latest_user_messages(request)
-
- # Nothing to do if the user_messages string is empty
- if len(user_messages) == 0:
+ # Get the latest user message
+ user_message = self.get_last_user_message_block(request)
+ if not user_message:
  return PipelineResult(request=request)
 
  # Create storage engine object
  storage_engine = StorageEngine()
 
  # Extract any code snippets
- snippets = extract_snippets(user_messages)
+ snippets = extract_snippets(user_message)
 
  bad_snippet_packages = []
  if len(snippets) > 0:
+ snippet_language = snippets[0].language
  # Collect all packages referenced in the snippets
  snippet_packages = []
  for snippet in snippets:
  snippet_packages.extend(
- PackageExtractor.extract_packages(snippet.code, snippet.language)
+ PackageExtractor.extract_packages(snippet.code, snippet.language) # type: ignore
  )
- logger.info(f"Found {len(snippet_packages)} packages in code snippets.")
 
+ logger.info(
+ f"Found {len(snippet_packages)} packages "
+ f"for language {snippet_language} in code snippets."
+ )
  # Find bad packages in the snippets
  bad_snippet_packages = await storage_engine.search(
- language=snippets[0].language, packages=snippet_packages
- )
+ language=snippet_language, packages=snippet_packages
+ ) # type: ignore
  logger.info(f"Found {len(bad_snippet_packages)} bad packages in code snippets.")
 
  # Remove code snippets from the user messages and search for bad packages
  # in the rest of the user query/messsages
- user_messages = re.sub(r"```.*?```", "", user_messages, flags=re.DOTALL)
-
- # Vector search to find bad packages
- bad_packages = await storage_engine.search(query=user_messages, distance=0.5, limit=100)
+ user_messages = re.sub(r"```.*?```", "", user_message, flags=re.DOTALL)
+ user_messages = re.sub(r"⋮...*?⋮...\n\n", "", user_messages, flags=re.DOTALL)
+
+ # split messages into double newlines, to avoid passing so many content in the search
+ split_messages = user_messages.split("\n\n")
+ collected_bad_packages = []
+ for item_message in split_messages:
+ # Vector search to find bad packages
+ bad_packages = await storage_engine.search(query=item_message, distance=0.5, limit=100)
+ if bad_packages and len(bad_packages) > 0:
+ collected_bad_packages.extend(bad_packages)
 
  # All bad packages
- all_bad_packages = bad_snippet_packages + bad_packages
+ all_bad_packages = bad_snippet_packages + collected_bad_packages
 
  logger.info(f"Adding {len(all_bad_packages)} bad packages to the context.")
 
@@ -119,7 +128,7 @@ async def process(
  # Add the context to the last user message
  # Format: "Context: {context_str} \n Query: {last user message content}"
  message = new_request["messages"][last_user_idx]
- context_msg = f'Context: {context_str} \n\n Query: {message["content"]}'
+ context_msg = f'Context: {context_str} \n\n Query: {message["content"]}' # type: ignore
  message["content"] = context_msg
 
  logger.debug("Final context message", context_message=context_msg)

src/codegate/pipeline/extract_snippets/extract_snippets.py

@@ -4,6 +4,7 @@
 
 import structlog
 from litellm.types.llms.openai import ChatCompletionRequest
+from pygments.lexers import guess_lexer
 
 from codegate.pipeline.base import CodeSnippet, PipelineContext, PipelineResult, PipelineStep
 
@@ -65,6 +66,8 @@ def ecosystem_from_message(message: str) -> Optional[str]:
  "ts": "typescript",
  "tsx": "typescript",
  "go": "go",
+ "rs": "rust",
+ "java": "java",
  }
  return language_mapping.get(message, None)
 
@@ -82,6 +85,7 @@ def extract_snippets(message: str) -> List[CodeSnippet]:
  # Regular expression to find code blocks
 
  snippets: List[CodeSnippet] = []
+ available_languages = ["python", "javascript", "typescript", "go", "rust", "java"]
 
  # Find all code block matches
  for match in CODE_BLOCK_PATTERN.finditer(message):
@@ -105,6 +109,14 @@ def extract_snippets(message: str) -> List[CodeSnippet]:
  filename = filename.strip()
  # Determine language from the filename
  lang = ecosystem_from_filepath(filename)
+ if lang is None:
+ # try to guess it from the code
+ lexer = guess_lexer(content)
+ if lexer and lexer.name:
+ lang = lexer.name.lower()
+ # only add available languages
+ if lang not in available_languages:
+ lang = None
 
  snippets.append(CodeSnippet(filepath=filename, code=content, language=lang))
 
@@ -129,10 +141,9 @@ async def process(
  request: ChatCompletionRequest,
  context: PipelineContext,
  ) -> PipelineResult:
-last_user_message = self.get_last_user_message(request)
- if not last_user_message:
+msg_content = self.get_last_user_message_block(request)
+ if not msg_content:
  return PipelineResult(request=request, context=context)
- msg_content, _ = last_user_message
  snippets = extract_snippets(msg_content)
 
  logger.info(f"Extracted {len(snippets)} code snippets from the user message")