Skip to content

Commit

Permalink
added llm for code detection too.
Browse files Browse the repository at this point in the history
  • Loading branch information
hrshdhgd committed Feb 10, 2024
1 parent 30047e2 commit 6c2fd00
Showing 1 changed file with 13 additions and 12 deletions.
25 changes: 13 additions & 12 deletions src/codergpt/main.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,5 @@
"""Main python file."""

import ast
import os
from pathlib import Path
from typing import Optional, Union
Expand All @@ -13,7 +12,6 @@
from codergpt.commenter.commenter import CodeCommenter
from codergpt.constants import EXTENSION_MAP_FILE, INSPECTION_HEADERS
from codergpt.explainer.explainer import CodeExplainer
from codergpt.utils.expression_evaluator import ExpressionEvaluator


class CoderGPT:
Expand Down Expand Up @@ -79,17 +77,20 @@ def get_code(

language_map = self.inspect_package(filename)
language = language_map.get(str(filename))

parsed_code = ast.parse(source_code)

visitor = ExpressionEvaluator(source_code=source_code, function_name=function_name, class_name=class_name)
visitor.visit(parsed_code)
if function_name:
return (visitor.function_code, language)
elif class_name:
return (visitor.class_code, language)
search_term = function_name if function_name else class_name
if search_term:
response = self.chain.invoke(
{
"input": f"Identify the structure of this {language} code \n{source_code}\n"
f" and give me only the code (no explanation) that corresponds"
f"to the {search_term} function or class."
}
)
code = response.content
else:
return (source_code, language)
code = source_code

return (code, language)

def explainer(self, path: Union[str, Path], function: str = None, classname=None):
"""
Expand Down

0 comments on commit 6c2fd00

Please sign in to comment.