clean up grpc calls

lunatrace/bsl/ingest-worker/pkg/ml/types.go

@@ -1,3 +1,14 @@
+// Copyright by LunaSec (owned by Refinery Labs, Inc)
+//
+// Licensed under the Business Source License v1.1 
+// (the "License"); you may not use this file except in compliance with the
+// License. You may obtain a copy of the License at
+//
+// https://github.com/lunasec-io/lunasec/blob/master/licenses/BSL-LunaTrace.txt
+//
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
 package ml
 
 import (
@@ -14,6 +25,8 @@ type ReferenceContent struct {
 	URL                 string
 	Content             string
 	NormalizedContent   string
+	ParsedContent       string
+	Summary             string
 	ContentType         string
 	LastSuccessfulFetch *time.Time
 }

lunatrace/bsl/ingest-worker/pkg/vulnerability/scrape/scraper.go

@@ -1,3 +1,14 @@
+// Copyright by LunaSec (owned by Refinery Labs, Inc)
+//
+// Licensed under the Business Source License v1.1 
+// (the "License"); you may not use this file except in compliance with the
+// License. You may obtain a copy of the License at
+//
+// https://github.com/lunasec-io/lunasec/blob/master/licenses/BSL-LunaTrace.txt
+//
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
 package scrape
 
 import (
@@ -86,14 +97,16 @@ func (p *scraper) scrapeVulnerabilityReference(ref *ReferenceInfo) *model.Refere
 	pRef.NormalizedContent = normalContent
 
 	// TODO (cthompson) store the cleaned content in the database
-	_, err = p.deps.LangChain.CleanWebpage(context.Background(), &gen.CleanWebpageRequest{
+	result, err := p.deps.LangChain.CleanAdvisory(context.Background(), &gen.CleanAdvisoryRequest{
 		Content:     normalContent,
 		Description: ref.VulnDesc,
 	})
 	if err != nil {
 		return &pRef
 	}
 
+	pRef.Summary = result.
+
 	return &pRef
 }
 

lunatrace/bsl/ml/python/README.md

@@ -1,4 +1,16 @@
-
+<!--
+  ~ Copyright by LunaSec (owned by Refinery Labs, Inc)
+  ~
+  ~ Licensed under the Creative Commons Attribution-ShareAlike 4.0 International
+  ~ (the "License"); you may not use this file except in compliance with the
+  ~ License. You may obtain a copy of the License at
+  ~
+  ~ https://creativecommons.org/licenses/by-sa/4.0/legalcode
+  ~
+  ~ See the License for the specific language governing permissions and
+  ~ limitations under the License.
+  ~
+-->
 # Lunasec-ml Python package
 
 This package contains lunasec's python ML tools, primarily built around langchain. 
@@ -45,8 +57,13 @@ The protobuf definitions are in `lunasec/lunatrace/bsl/proto/`. To update the pr
 yay -S protobuf
 python -m pip install grpcio-tools
 go install google.golang.org/protobuf/cmd/protoc-gen-go@latest
+go install google.golang.org/grpc/cmd/protoc-gen-go-grpc@latest
 ```
 
+To regenerate all generated grpc client code, run this from the repo root
+```shell
+yarn run generate:go && yarn run generate:ml:python
+```
 To regenerate the grpc python server, run:
 ```
 python -m grpc_tools.protoc -I../../proto --python_out=. --pyi_out=. --grpc_python_out=. ../../proto/langchain.proto

lunatrace/bsl/ml/python/client.py

@@ -1,10 +1,11 @@
+# TODO: IS THIS DEAD CODE IN THIS FILE?
+
 from twirp.context import Context
 from twirp.exceptions import TwirpServerException
 
 from gen import langchain_twirp, langchain_pb2
 
 client = langchain_twirp.LangChainClient("http://localhost:3000")
-
 # if you are using a custom prefix, then pass it as `server_path_prefix`
 # param to `MakeHat` class.
 try:

lunatrace/bsl/ml/python/langchain_pb2.pyi

@@ -1,6 +1,7 @@
+from google.protobuf.internal import containers as _containers
 from google.protobuf import descriptor as _descriptor
 from google.protobuf import message as _message
-from typing import ClassVar as _ClassVar, Optional as _Optional
+from typing import ClassVar as _ClassVar, Iterable as _Iterable, Mapping as _Mapping, Optional as _Optional, Union as _Union
 
 DESCRIPTOR: _descriptor.FileDescriptor
 
@@ -11,35 +12,61 @@ class ChatRequest(_message.Message):
     def __init__(self, message: _Optional[str] = ...) -> None: ...
 
 class ChatResponse(_message.Message):
-    __slots__ = ["response"]
-    RESPONSE_FIELD_NUMBER: _ClassVar[int]
-    response: str
-    def __init__(self, response: _Optional[str] = ...) -> None: ...
+    __slots__ = ["finalAnswer", "intermediateSteps"]
+    FINALANSWER_FIELD_NUMBER: _ClassVar[int]
+    INTERMEDIATESTEPS_FIELD_NUMBER: _ClassVar[int]
+    finalAnswer: str
+    intermediateSteps: str
+    def __init__(self, finalAnswer: _Optional[str] = ..., intermediateSteps: _Optional[str] = ...) -> None: ...
 
-class CleanWebpageRequest(_message.Message):
+class CleanAdvisoryRequest(_message.Message):
     __slots__ = ["content", "description"]
     CONTENT_FIELD_NUMBER: _ClassVar[int]
     DESCRIPTION_FIELD_NUMBER: _ClassVar[int]
     content: str
     description: str
     def __init__(self, content: _Optional[str] = ..., description: _Optional[str] = ...) -> None: ...
 
-class CleanWebpageResponse(_message.Message):
-    __slots__ = ["content"]
+class CleanAdvisoryResponse(_message.Message):
+    __slots__ = ["content", "description"]
     CONTENT_FIELD_NUMBER: _ClassVar[int]
+    DESCRIPTION_FIELD_NUMBER: _ClassVar[int]
     content: str
-    def __init__(self, content: _Optional[str] = ...) -> None: ...
+    description: str
+    def __init__(self, content: _Optional[str] = ..., description: _Optional[str] = ...) -> None: ...
 
-class SummarizeRequest(_message.Message):
-    __slots__ = ["content", "query"]
-    CONTENT_FIELD_NUMBER: _ClassVar[int]
-    QUERY_FIELD_NUMBER: _ClassVar[int]
-    content: str
-    query: str
-    def __init__(self, content: _Optional[str] = ..., query: _Optional[str] = ...) -> None: ...
+class CleanSnippetsRequest(_message.Message):
+    __slots__ = ["snippets"]
+    SNIPPETS_FIELD_NUMBER: _ClassVar[int]
+    snippets: _containers.RepeatedCompositeFieldContainer[Snippet]
+    def __init__(self, snippets: _Optional[_Iterable[_Union[Snippet, _Mapping]]] = ...) -> None: ...
 
-class SummarizeResponse(_message.Message):
-    __slots__ = ["summary"]
+class CleanSnippetsResponse(_message.Message):
+    __slots__ = ["snippets"]
+    SNIPPETS_FIELD_NUMBER: _ClassVar[int]
+    snippets: _containers.RepeatedCompositeFieldContainer[CleanedSnippet]
+    def __init__(self, snippets: _Optional[_Iterable[_Union[CleanedSnippet, _Mapping]]] = ...) -> None: ...
+
+class CleanedSnippet(_message.Message):
+    __slots__ = ["code", "language", "score", "summary", "type"]
+    CODE_FIELD_NUMBER: _ClassVar[int]
+    LANGUAGE_FIELD_NUMBER: _ClassVar[int]
+    SCORE_FIELD_NUMBER: _ClassVar[int]
     SUMMARY_FIELD_NUMBER: _ClassVar[int]
+    TYPE_FIELD_NUMBER: _ClassVar[int]
+    code: str
+    language: str
+    score: int
     summary: str
-    def __init__(self, summary: _Optional[str] = ...) -> None: ...
+    type: str
+    def __init__(self, code: _Optional[str] = ..., score: _Optional[int] = ..., summary: _Optional[str] = ..., type: _Optional[str] = ..., language: _Optional[str] = ...) -> None: ...
+
+class Snippet(_message.Message):
+    __slots__ = ["code", "preamble", "vuln_description"]
+    CODE_FIELD_NUMBER: _ClassVar[int]
+    PREAMBLE_FIELD_NUMBER: _ClassVar[int]
+    VULN_DESCRIPTION_FIELD_NUMBER: _ClassVar[int]
+    code: str
+    preamble: str
+    vuln_description: str
+    def __init__(self, code: _Optional[str] = ..., preamble: _Optional[str] = ..., vuln_description: _Optional[str] = ...) -> None: ...

lunatrace/bsl/ml/python/langchain_pb2_grpc.py

@@ -14,39 +14,40 @@ def __init__(self, channel):
         Args:
             channel: A grpc.Channel.
         """
-        self.Summarize = channel.unary_unary(
-                '/langchain.LangChain/Summarize',
-                request_serializer=langchain__pb2.SummarizeRequest.SerializeToString,
-                response_deserializer=langchain__pb2.SummarizeResponse.FromString,
-                )
-        self.CleanWebpage = channel.unary_unary(
-                '/langchain.LangChain/CleanWebpage',
-                request_serializer=langchain__pb2.CleanWebpageRequest.SerializeToString,
-                response_deserializer=langchain__pb2.CleanWebpageResponse.FromString,
+        self.CleanAdvisory = channel.unary_unary(
+                '/langchain.LangChain/CleanAdvisory',
+                request_serializer=langchain__pb2.CleanAdvisoryRequest.SerializeToString,
+                response_deserializer=langchain__pb2.CleanAdvisoryResponse.FromString,
                 )
         self.Chat = channel.unary_unary(
                 '/langchain.LangChain/Chat',
                 request_serializer=langchain__pb2.ChatRequest.SerializeToString,
                 response_deserializer=langchain__pb2.ChatResponse.FromString,
                 )
+        self.CleanSnippets = channel.unary_unary(
+                '/langchain.LangChain/CleanSnippets',
+                request_serializer=langchain__pb2.CleanSnippetsRequest.SerializeToString,
+                response_deserializer=langchain__pb2.CleanSnippetsResponse.FromString,
+                )
 
 
 class LangChainServicer(object):
     """Missing associated documentation comment in .proto file."""
 
-    def Summarize(self, request, context):
-        """Missing associated documentation comment in .proto file."""
+    def CleanAdvisory(self, request, context):
+        """rpc Summarize(SummarizeRequest) returns (SummarizeResponse);
+        """
         context.set_code(grpc.StatusCode.UNIMPLEMENTED)
         context.set_details('Method not implemented!')
         raise NotImplementedError('Method not implemented!')
 
-    def CleanWebpage(self, request, context):
+    def Chat(self, request, context):
         """Missing associated documentation comment in .proto file."""
         context.set_code(grpc.StatusCode.UNIMPLEMENTED)
         context.set_details('Method not implemented!')
         raise NotImplementedError('Method not implemented!')
 
-    def Chat(self, request, context):
+    def CleanSnippets(self, request, context):
         """Missing associated documentation comment in .proto file."""
         context.set_code(grpc.StatusCode.UNIMPLEMENTED)
         context.set_details('Method not implemented!')
@@ -55,21 +56,21 @@ def Chat(self, request, context):
 
 def add_LangChainServicer_to_server(servicer, server):
     rpc_method_handlers = {
-            'Summarize': grpc.unary_unary_rpc_method_handler(
-                    servicer.Summarize,
-                    request_deserializer=langchain__pb2.SummarizeRequest.FromString,
-                    response_serializer=langchain__pb2.SummarizeResponse.SerializeToString,
-            ),
-            'CleanWebpage': grpc.unary_unary_rpc_method_handler(
-                    servicer.CleanWebpage,
-                    request_deserializer=langchain__pb2.CleanWebpageRequest.FromString,
-                    response_serializer=langchain__pb2.CleanWebpageResponse.SerializeToString,
+            'CleanAdvisory': grpc.unary_unary_rpc_method_handler(
+                    servicer.CleanAdvisory,
+                    request_deserializer=langchain__pb2.CleanAdvisoryRequest.FromString,
+                    response_serializer=langchain__pb2.CleanAdvisoryResponse.SerializeToString,
             ),
             'Chat': grpc.unary_unary_rpc_method_handler(
                     servicer.Chat,
                     request_deserializer=langchain__pb2.ChatRequest.FromString,
                     response_serializer=langchain__pb2.ChatResponse.SerializeToString,
             ),
+            'CleanSnippets': grpc.unary_unary_rpc_method_handler(
+                    servicer.CleanSnippets,
+                    request_deserializer=langchain__pb2.CleanSnippetsRequest.FromString,
+                    response_serializer=langchain__pb2.CleanSnippetsResponse.SerializeToString,
+            ),
     }
     generic_handler = grpc.method_handlers_generic_handler(
             'langchain.LangChain', rpc_method_handlers)
@@ -81,7 +82,7 @@ class LangChain(object):
     """Missing associated documentation comment in .proto file."""
 
     @staticmethod
-    def Summarize(request,
+    def CleanAdvisory(request,
             target,
             options=(),
             channel_credentials=None,
@@ -91,14 +92,14 @@ def Summarize(request,
             wait_for_ready=None,
             timeout=None,
             metadata=None):
-        return grpc.experimental.unary_unary(request, target, '/langchain.LangChain/Summarize',
-            langchain__pb2.SummarizeRequest.SerializeToString,
-            langchain__pb2.SummarizeResponse.FromString,
+        return grpc.experimental.unary_unary(request, target, '/langchain.LangChain/CleanAdvisory',
+            langchain__pb2.CleanAdvisoryRequest.SerializeToString,
+            langchain__pb2.CleanAdvisoryResponse.FromString,
             options, channel_credentials,
             insecure, call_credentials, compression, wait_for_ready, timeout, metadata)
 
     @staticmethod
-    def CleanWebpage(request,
+    def Chat(request,
             target,
             options=(),
             channel_credentials=None,
@@ -108,14 +109,14 @@ def CleanWebpage(request,
             wait_for_ready=None,
             timeout=None,
             metadata=None):
-        return grpc.experimental.unary_unary(request, target, '/langchain.LangChain/CleanWebpage',
-            langchain__pb2.CleanWebpageRequest.SerializeToString,
-            langchain__pb2.CleanWebpageResponse.FromString,
+        return grpc.experimental.unary_unary(request, target, '/langchain.LangChain/Chat',
+            langchain__pb2.ChatRequest.SerializeToString,
+            langchain__pb2.ChatResponse.FromString,
             options, channel_credentials,
             insecure, call_credentials, compression, wait_for_ready, timeout, metadata)
 
     @staticmethod
-    def Chat(request,
+    def CleanSnippets(request,
             target,
             options=(),
             channel_credentials=None,
@@ -125,8 +126,8 @@ def Chat(request,
             wait_for_ready=None,
             timeout=None,
             metadata=None):
-        return grpc.experimental.unary_unary(request, target, '/langchain.LangChain/Chat',
-            langchain__pb2.ChatRequest.SerializeToString,
-            langchain__pb2.ChatResponse.FromString,
+        return grpc.experimental.unary_unary(request, target, '/langchain.LangChain/CleanSnippets',
+            langchain__pb2.CleanSnippetsRequest.SerializeToString,
+            langchain__pb2.CleanSnippetsResponse.FromString,
             options, channel_credentials,
             insecure, call_credentials, compression, wait_for_ready, timeout, metadata)