fix problems with snippet language detection and search

Author: yrobla

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"}

poetry.lock

@@ -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"

pyproject.toml

@@ -1,5 +1,7 @@
 import json
 import re
+import malicious_pypi_dummy
+
 
 import structlog
 from litellm import ChatCompletionRequest
@@ -72,18 +74,17 @@ async def process(
 
         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)  # type: ignore
                 )
-            logger.info(f"Found {len(snippet_packages)} packages in code snippets.")
 
+            logger.info(f"Found {len(snippet_packages)} packages 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  # type: ignore
-            )
+            bad_snippet_packages = await storage_engine.search(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

src/codegate/pipeline/codegate_context_retriever/codegate.py

@@ -1,5 +1,6 @@
 import os
 import re
+from pygments.lexers import guess_lexer
 from typing import List, Optional
 
 import structlog
@@ -85,7 +86,6 @@ def extract_snippets(message: str) -> List[CodeSnippet]:
 
     # Find all code block matches
     for match in CODE_BLOCK_PATTERN.finditer(message):
-        print("i try to extract snippet")
         matched_language = match.group("language") if match.group("language") else None
         filename = match.group("filename") if match.group("filename") else None
         content = match.group("content")
@@ -95,24 +95,23 @@ def extract_snippets(message: str) -> List[CodeSnippet]:
         # format ` ```python ` in output snippets
         if filename and not matched_language and "." not in filename:
             lang = filename
-            print("lang is")
-            print(lang)
             filename = None
         else:
             # Determine language from the message, either by the short
             # language identifier or by the filename
             lang = None
             if matched_language:
-                print("i have a matched language")
                 lang = ecosystem_from_message(matched_language.strip())
             if lang is None and filename:
-                print("I try to get from filename")
                 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()
 
-        print("language is")
-        print(lang)
         snippets.append(CodeSnippet(filepath=filename, code=content, language=lang))
 
     return snippets

src/codegate/pipeline/extract_snippets/extract_snippets.py

@@ -74,21 +74,16 @@ class PackageExtractor:
 
     @staticmethod
     def extract_packages(code: str, language_name: str) -> list[str]:
-        print("packages are")
-        print(code)
-        print(language_name)
         if (code is None) or (language_name is None):
             return []
 
         language_name = language_name.lower()
 
         if language_name not in PackageExtractor.__languages.keys():
-            print("no langauge")
             return []
 
         language = PackageExtractor.__languages[language_name]
         parser = PackageExtractor.__parsers[language_name]
-        print("here")
 
         # Create tree
         tree = parser.parse(bytes(code, "utf8"))