Merge pull request #202 from unoplat/108-feat-enable-function-call-li…

…nkage-with-content-also-enable-all-possible-function-metadata

108 feat enable function call linkage with content also enable all possible function metadata

unoplat-code-confluence-commons/unoplat_code_confluence_commons/graph_models/base_models.py

@@ -13,3 +13,8 @@ class BaseNode(StructuredNode):
 class ContainsRelationship(StructuredRel):
     """Relationship for representing containment between nodes"""
     pass
+
+class CallsRelationship(StructuredRel):
+    """Represents a method call from one method to another."""
+    parameters = JSONProperty()
+    position = JSONProperty()

unoplat-code-confluence-commons/unoplat_code_confluence_commons/graph_models/confluence_annotation.py

@@ -0,0 +1,9 @@
+from neomodel import StructuredNode, StringProperty, RelationshipFrom, ZeroOrMore, JSONProperty
+
+class ConfluenceAnnotation(StructuredNode):
+    name = StringProperty(required=True)
+    key_values = JSONProperty()
+    position = JSONProperty()
+    # Relationships
+    annotated_classes = RelationshipFrom('.confluence_class.ConfluenceClass', 'HAS_ANNOTATION', cardinality=ZeroOrMore)
+    annotated_methods = RelationshipFrom('.confluence_method.ConfluenceMethod', 'HAS_ANNOTATION', cardinality=ZeroOrMore)

unoplat-code-confluence-commons/unoplat_code_confluence_commons/graph_models/confluence_class.py

@@ -1,5 +1,5 @@
 from .base_models import BaseNode, ContainsRelationship
-from neomodel import RelationshipFrom, RelationshipTo, StringProperty,ZeroOrMore,One,ArrayProperty,VectorIndex,FloatProperty
+from neomodel import RelationshipFrom, RelationshipTo, StringProperty,ZeroOrMore,One,ArrayProperty,VectorIndex,FloatProperty,JSONProperty
 
 class ConfluenceClass(BaseNode):
     """Represents a class in a package"""
@@ -8,6 +8,17 @@ class ConfluenceClass(BaseNode):
     objective = StringProperty(default="")
     class_objective_embedding = ArrayProperty(FloatProperty())
     class_implementation_summary_embedding = ArrayProperty(FloatProperty())
+    node_type = StringProperty()
+    file_path = StringProperty()
+    module = StringProperty()
+    extend = StringProperty()
+    multiple_extend = ArrayProperty(StringProperty())
+    position = JSONProperty()
+    content = StringProperty()
     # Class relationships
     package = RelationshipTo('.confluence_package.ConfluencePackage', 'BELONGS_TO', model=ContainsRelationship, cardinality=One)
     methods = RelationshipTo('.confluence_method.ConfluenceMethod', 'CONTAINS', model=ContainsRelationship, cardinality=ZeroOrMore)
+    extends = RelationshipTo('.confluence_class.ConfluenceClass', 'EXTENDS', cardinality=ZeroOrMore)
+    imports = RelationshipTo('.confluence_import.ConfluenceImport', 'IMPORTS', cardinality=ZeroOrMore)
+    annotations = RelationshipTo('.confluence_annotation.ConfluenceAnnotation', 'HAS_ANNOTATION', cardinality=ZeroOrMore)
+    fields = RelationshipTo('.confluence_class_field.ConfluenceClassField', 'CONTAINS', model=ContainsRelationship, cardinality=ZeroOrMore)

unoplat-code-confluence-commons/unoplat_code_confluence_commons/graph_models/confluence_class_field.py

@@ -0,0 +1,6 @@
+from neomodel import StructuredNode, StringProperty, RelationshipTo, ZeroOrMore
+
+class ConfluenceClassField(StructuredNode):
+    field_type = StringProperty()
+    field_name = StringProperty()
+    annotations = RelationshipTo('.confluence_annotation.ConfluenceAnnotation', 'HAS_ANNOTATION', cardinality=ZeroOrMore)

unoplat-code-confluence-commons/unoplat_code_confluence_commons/graph_models/confluence_import.py

@@ -0,0 +1,8 @@
+from neomodel import StructuredNode, StringProperty, ArrayProperty, RelationshipFrom, ZeroOrMore
+
+class ConfluenceImport(StructuredNode):
+    source = StringProperty(required=True)
+
+    usage_names = ArrayProperty(StringProperty())
+
+    imported_by = RelationshipFrom('.confluence_class.ConfluenceClass', 'IMPORTS', cardinality=ZeroOrMore)

unoplat-code-confluence-commons/unoplat_code_confluence_commons/graph_models/confluence_method.py

@@ -1,12 +1,21 @@
-from .base_models import BaseNode, ContainsRelationship
-from neomodel import RelationshipTo, StringProperty,One,ArrayProperty,FloatProperty
+from .base_models import BaseNode, ContainsRelationship, CallsRelationship
+from neomodel import RelationshipTo, StringProperty,One,ArrayProperty,FloatProperty,ZeroOrMore,IntegerProperty,JSONProperty
 
 class ConfluenceMethod(BaseNode):
     """Represents a method in a class"""
+
     function_name = StringProperty(required=True)
+    return_type = StringProperty()
     implementation_summary = StringProperty(default="")
     objective = StringProperty(default="")
     function_objective_embedding = ArrayProperty(FloatProperty())
     function_implementation_summary_embedding = ArrayProperty(FloatProperty())
+    content = StringProperty()
+    body_hash = IntegerProperty()
+    position = JSONProperty()
+    local_variables = JSONProperty()
+
     # Method relationships
     confluence_class = RelationshipTo('.confluence_class.ConfluenceClass', 'BELONGS_TO', model=ContainsRelationship, cardinality=One)
+    annotations = RelationshipTo('.confluence_annotation.ConfluenceAnnotation', 'HAS_ANNOTATION', cardinality=ZeroOrMore)
+    calls_methods = RelationshipTo('.confluence_method.ConfluenceMethod', 'CALLS', model=CallsRelationship, cardinality=ZeroOrMore)

unoplat-code-confluence/unoplat_code_confluence/__main__.py

@@ -128,7 +128,7 @@ async def start_parsing(app_config: AppConfig, iload_json: JsonLoader, iparse_js
 
     unoplat_codebase_summary.codebase_name = app_config.codebase_name
 
-    unoplat_graph_processing.process_codebase_summary(unoplat_codebase_summary)
+    unoplat_graph_processing.process_codebase_summary(unoplat_codebase,unoplat_codebase_summary)
 
     markdown_output = isummariser.summarise_to_markdown(unoplat_codebase_summary)
     # write the markdown output to a file

unoplat-code-confluence/unoplat_code_confluence/configuration/external_config.py

@@ -3,13 +3,7 @@
 from pydantic import BaseModel, ValidationInfo, field_validator, Field
 
 
-class LLMProvider(Enum):
-    OPENAI = 'openai'
-    COHERE = 'cohere'
-    ANYSCALE = 'anyscale'
-    TOGETHER = 'together'
-    OLLAMA = 'ollama'
-    AWSANTHROPIC = 'awsanthropic'
+
 
 
 
@@ -37,7 +31,17 @@ class AppConfig(BaseModel):
     programming_language: str
     repo: RepoConfig
     api_tokens: Dict[str, str]
-    llm_provider_config: Dict[str, Any]
+    llm_provider_config: Dict[str, Any] = Field(
+        description="Configuration for the LLM provider based on litellm",
+        example={
+            "model_provider": "openai/gpt-4",
+            "model_provider_args": {
+                "api_key": "sk-...",
+                "max_tokens": 500,
+                "temperature": 0.0
+            }
+        }
+    )
     handlers: List[Dict[str, Any]] = Field(default_factory=list,alias="logging_handlers")
     parallisation: int = 1
     json_output: bool = False
@@ -60,10 +64,3 @@ def check_api_tokens(cls, value, info:ValidationInfo):
         if len(value) != 1:
             raise ValueError("api_tokens must only contain github_token")
         return value
-
-    @field_validator('llm_provider_config')
-    def check_llm_provider_config(cls, value, info:ValidationInfo):
-        #TODO if key is in LLmProvider enum
-        if not all(key in LLMProvider for key in value.keys()):
-            raise ValueError("llm_provider_config keys must be in LLMProvider enum")
-        return value

unoplat-code-confluence/unoplat_code_confluence/data_models/chapi_unoplat_function.py

@@ -6,12 +6,12 @@
 from unoplat_code_confluence.data_models.chapi_unoplat_annotation import Annotation
 from unoplat_code_confluence.data_models.chapi_unoplat_position import Position
 
-class Function(BaseModel):
+class ChapiUnoplatFunction(BaseModel):
     name: Optional[str] = Field(default=None, alias="Name")
     return_type: Optional[str] = Field(default=None, alias="ReturnType")
     function_calls: List[FunctionCall] = Field(default_factory=list, alias="FunctionCalls")
     annotations: List[Annotation] = Field(default_factory=list, alias="Annotations")
-    position: Optional[Position] = Field(default=None, alias="Position")
+    position: Optional[Position] = Field(default=None, alias="Position",exclude=True)
     local_variables: List[UnoplatFunctionFieldModel] = Field(default_factory=list, alias="LocalVariables")
-    body_hash: Optional[int] = Field(default=None, alias="BodyHash")
+    body_hash: Optional[int] = Field(default=None, alias="BodyHash",exclude=True)
     content: Optional[str] = Field(default=None, alias="Content")

unoplat-code-confluence/unoplat_code_confluence/data_models/chapi_unoplat_functioncall.py

@@ -9,9 +9,13 @@
 
 
 class FunctionCall(BaseModel):
-    package: Optional[str] = Field(default=None, alias="Package")
-    type: Optional[str] = Field(default=None, alias="Type")
-    node_name: Optional[str] = Field(default=None, alias="NodeName")
-    function_name: Optional[str] = Field(default=None, alias="FunctionName")
-    parameters: List[Parameter] = Field(default_factory=list, alias="Parameters")
-    position: Optional[Position] = Field(default=None, alias="Position")
+    """
+    FunctionCall is a data model for a function call being made in a function.
+    """
+    package: Optional[str] = Field(default=None, alias="Package",description="package name of the function call")
+    type: Optional[str] = Field(default=None, alias="Type",description="type of the function call")
+    node_name: Optional[str] = Field(default=None, alias="NodeName",description="name of the node being called")
+    function_name: Optional[str] = Field(default=None, alias="FunctionName",description="name of the function being called")
+    parameters: List[Parameter] = Field(default_factory=list, alias="Parameters",description="parameters of the function call")
+    position: Optional[Position] = Field(default=None, alias="Position",exclude=True,description="position of the function call in the code")
+

unoplat-code-confluence/unoplat_code_confluence/data_models/chapi_unoplat_node.py

@@ -3,22 +3,22 @@
 from unoplat_code_confluence.data_models.chapi_unoplat_annotation import Annotation
 from unoplat_code_confluence.data_models.chapi_unoplat_class_fieldmodel import ClassFieldModel
 from unoplat_code_confluence.data_models.chapi_unoplat_import import Import
-from unoplat_code_confluence.data_models.chapi_unoplat_function import Function
+from unoplat_code_confluence.data_models.chapi_unoplat_function import ChapiUnoplatFunction
 from unoplat_code_confluence.data_models.chapi_unoplat_position import Position
 
 
 
-class Node(BaseModel):
+class ChapiUnoplatNode(BaseModel):
     node_name: Optional[str] = Field(default=None, alias="NodeName")
     type: Optional[str] = Field(default=None, alias="Type")
-    file_path: Optional[str] = Field(default=None, alias="FilePath")
-    module: Optional[str] = Field(default=None, alias="Module")
-    package: Optional[str] = Field(default=None, alias="Package")
+    file_path: Optional[str] = Field(default=None, alias="FilePath",exclude=True)
+    module: Optional[str] = Field(default=None, alias="Module",exclude=True)
+    package: Optional[str] = Field(default=None, alias="Package",exclude=True)
     multiple_extend: Optional[list[str]] = Field(default_factory=list, alias="MultipleExtend")
     fields: List[ClassFieldModel] = Field(default_factory=list, alias="Fields")
     extend: Optional[str] = Field(default=None, alias="Extend")
     imports: List[Import] = Field(default_factory=list, alias="Imports")
-    functions: List[Function] = Field(default_factory=list, alias="Functions")
-    position: Optional[Position] = Field(default=None, alias="Position")
-    content: Optional[str] = Field(default=None, alias="Content")
+    functions: List[ChapiUnoplatFunction] = Field(default_factory=list, alias="Functions")
+    position: Optional[Position] = Field(default=None, alias="Position",exclude=True)
+    content: Optional[str] = Field(default=None, alias="Content",exclude=True)
     annotations: List[Annotation] = Field(default_factory=list, alias="Annotations")

unoplat-code-confluence/unoplat_code_confluence/data_models/chapi_unoplat_package.py

@@ -1,12 +1,12 @@
 from typing import Dict, List, Optional
 from pydantic import BaseModel, Field
 
-from unoplat_code_confluence.data_models.dspy.dspy_unoplat_fs_node_subset import DspyUnoplatNodeSubset
+from unoplat_code_confluence.data_models.chapi_unoplat_node import ChapiUnoplatNode
 
 
 class UnoplatPackage(BaseModel):
     name: Optional[str] = Field(default=None,description="Name of the package")
-    node_subsets: Optional[List[DspyUnoplatNodeSubset]] = Field( default_factory=list,description="List of the node subsets for the package")
+    nodes: Optional[List[ChapiUnoplatNode]] = Field( default_factory=list,description="List of the nodes for the package")
     sub_packages: Optional[Dict[str, 'UnoplatPackage']] = Field( default_factory=dict,description="Dict of the sub-packages for the package")
 
 UnoplatPackage.model_rebuild()

unoplat-code-confluence/unoplat_code_confluence/data_models/dspy/dspy_unoplat_function_summary.py

@@ -1,9 +1,9 @@
 from typing import Optional
 from pydantic import BaseModel,Field
 
-from unoplat_code_confluence.data_models.dspy.dspy_o_function_summary import DspyFunctionSummary
 
 class DspyUnoplatFunctionSummary(BaseModel):
     function_name: str = Field( alias="FunctionName", description="The name of the function")
-    function_summary: DspyFunctionSummary = Field( alias="FunctionSummary", description="A summary of the function")
+    objective: Optional[str] = Field(default=None, alias="Objective",description="This should include high level objective of what function does based on function content and function metadata. Should not be more than 3 lines.")
+    implementation_summary: Optional[str] = Field(default=None, alias="ImplementationSummary",description="This should include implementation details of the function. make sure if this function makes internal calls to other functions of same class and to external calls to other classes/libs is also covered. Use all metadata shared for the function to answer .")
     metadata: Optional[dict] = Field(default=None, description="Additional metadata for the function")

unoplat-code-confluence/unoplat_code_confluence/database/graph/unoplat_graph_ingestion.py

@@ -8,7 +8,6 @@ def __init__(self, app_config: AppConfig):
         uri = app_config.neo4j_uri
         username = app_config.neo4j_username
         password = app_config.neo4j_password
-        #self.driver = GraphDatabase.driver(uri, auth=(username, password))
         config.DATABASE_URL = f'bolt://{username}:{password}@{uri.split("://")[-1]}'
         db.set_connection(config.DATABASE_URL)
 

unoplat-code-confluence/unoplat_code_confluence/database/graph/unoplat_graph_processing.py

@@ -1,3 +1,4 @@
+from unoplat_code_confluence.data_models.chapi_unoplat_codebase import UnoplatCodebase
 from unoplat_code_confluence.data_models.dspy.dspy_unoplat_function_summary import DspyUnoplatFunctionSummary
 from unoplat_code_confluence.data_models.dspy.dspy_unoplat_node_summary import DspyUnoplatNodeSummary
 from unoplat_code_confluence.data_models.dspy.dspy_unoplat_package_summary import DspyUnoplatPackageSummary
@@ -22,7 +23,7 @@ def __init__(self, app_config: AppConfig):
         db.set_connection(neo4j_url)
         self.unoplat_graph_ingestion.create_schema()
 
-    def process_codebase_summary(self, codebase_summary: DspyUnoplatCodebaseSummary):
+    def process_codebase_summary(self, unoplat_codebase: UnoplatCodebase, codebase_summary: DspyUnoplatCodebaseSummary):
         try:
             # Generate embeddings
             objective_embedding = self.embedding_generator.generate_embeddings_for_single_text(
@@ -166,9 +167,9 @@ def _process_method(self, class_node: ConfluenceClass, method_summary: DspyUnopl
             method_node_dict = [{
                 "qualified_name": f"{class_node.qualified_name}.{method_summary.function_name}",
                 "function_name": method_summary.function_name,
-                "objective": method_summary.function_summary.objective,
+                "objective": method_summary.objective,
                 "function_objective_embedding": method_objective_embedding,
-                "implementation_summary": method_summary.function_summary.implementation_summary,
+                "implementation_summary": method_summary.implementation_summary,
                 "function_implementation_summary_embedding": method_implementation_embedding
             }]