Coverage for src/configuration/config_algo.py: 48%

1"""The main driver for the CueCode configuration algorithm""" 

2 

3from uuid import UUID 

4 

5from common.database_engine import DBEngine 

6from common.models.openapi_spec import OpenAPISpec 

7from configuration.openapi_schema_adapter import OpenAPISchemaAdapter 

8from configuration.openapi_schema_validate import validate_openapi_spec 

9from configuration.openapi_spec_entity_collection import OpenAPISpecEntityCollection 

10from configuration.openapi_validator_to_cuecode import ( 

11 create_selection_embeddings, 

12 openapi_spec_validator_to_cuecode_config, 

13) 

14 

15from .openapi import OpenAPIObject 

16 

17db_engine_for_workaround = DBEngine() # would be managed in dramatiq actor code 

18# if workaround to avoid Dramatiq fails. 

19 

20 

21def config_algo_openapi(db_engine: DBEngine, openapi_spec_id: str): 

22 """Runs the CueCode config algo against an OpenAPI spec. This is not the 

23 Dramatiq actor! If we need to use an actor, define a function that calls this 

24 one and then decorate the new function as the Dramatiq actor, all under a 

25 new 'actors' module.""" 

26 

27 db_engine = db_engine_for_workaround 

28 session = db_engine.get_session() 

29 

30 # Fetch OpenAPI spec from PostgreSQL 

31 db_spec = session.get(OpenAPISpec, openapi_spec_id) 

32 

33 # Ensure the text spec is a valid OpenAPI specification as such, apart any 

34 # knowledge of CueCode's requirements for OpenAPI spec structure. 

35 validate_openapi_spec(db_spec) 

36 

37 # Prepare the spec text for CueCode's parsing and validation, since 

38 # CueCode constrains some OpenAPI options and also provides extensions 

39 # to the OpenAPI spec 

40 spec_adapter = OpenAPISchemaAdapter(db_spec) 

41 spec_json = spec_adapter.get_cleaned_json_dict() 

42 

43 # Parse the OpenAPI specification 

44 parsed_spec = from_formatted_json(UUID(openapi_spec_id), spec_json) 

45 

46 # Pull from the parsed spec all SQLAlchemy entities represented in the spec 

47 # pylint: disable-next=unused-variable 

48 spec_entities: OpenAPISpecEntityCollection = ( 

49 openapi_spec_validator_to_cuecode_config(session, parsed_spec, db_spec) 

50 ) 

51 

52 session.add(db_spec) 

53 

54 create_selection_embeddings(db_spec, session) 

55 session.commit() 

56 

57 # NOTE The comments below describe the config algo from the Activity Diagram 

58 # in our Design docs, but it does NOT really follow how the code will be 

59 # structured if using OOP! A TODO is to clean up these comments in prep for 

60 # using the output of validator_to_entity_collection(). 

61 

62 # parallel over Schema Object 

63 

64 # parallel over HTTP verb in Path OpenAPI spec object 

65 

66 # Schema object has x-cuecode-exclude? Yes then 

67 # next item in loop 

68 # Initialize new OpenApiEntity object 

69 

70 # Schema object has x-cuecode-prompt? Yes then 

71 # Save x-cuecode-prompt to OpenApiEntity .nounPrompt 

72 # else 

73 # Save Schema Object name to OpenApiEntity .nounPrompt 

74 

75 # Call Ollama for embedding of nounPrompt 

76 

77 # parallel over Schema Object in OpenAPI spec 

78 

79 # UPSERT entity list to PostgreSQL (but do not commit transactinon) 

80 

81 # Call Ollama for embedding of selectionPrompt 

82 

83 # UPSERT endpoint list to PostgreSQL (but do not commit transactinon) 

84 

85 # Build graph of entity relationships based on $ref pointers in OpenAPI spec 

86 

87 # UPSERT graph edges to PostgreSQL as openapi_entity_dependency rows 

88 

89 # Encode OpenAPI Paths and verbs as LLM Tool calls with the .selectionPrompt chosen earlier 

90 

91 # Build LLM tool call JSON from Path objects and to PostgreSQL 

92 # do it 

93 

94 # Update status of config job in DB (user-facing record, not Dramatiq job) 

95 

96 # Commit PostgreSQL transaction 

97 

98 # Ack config job queue task 

99 

100 # print(openapi_spec) 

101 

102 # BEGIN Chase's work 

103 

104 # Needed in the event no server is specified in the servers array 

105 

106 

107def from_formatted_json(spec_id: UUID, data: dict): 

108 """create openapi object from json""" 

109 return OpenAPIObject( 

110 openapi_spec_uuid=spec_id, base_url=data["servers"][0]["url"], **data 

111 )