feat(wren-ai-service): adjust relation recommendation pipeline prompt and output spec

[paopa]

This PR updates the relationship recommendation pipeline. The changes focus on improving the pipeline's functionality, observability, and output structure.

Key Updates:

1. Added Observability:

   • Introduced @observe decorators to key functions for improved monitoring and debugging.

2. New cleaned_models Function:

   • Added a function to filter out existing relationships from the input models.

3. Updated Prompt Generation:

   • Modified the prompt function to use the cleaned models.

4. Enhanced System Prompt:

   • Significantly updated the system prompt to provide more specific guidelines for relationship recommendations.
   • Changed the output structure to include more detailed relationship information.

5. Updated ModelRelationship Class:

   • Revised the structure to match the new relationship recommendation format.

Detailed Changes

New cleaned_models Function:

@observe(capture_input=False)
def cleaned_models(mdl: dict) -> dict:
    def column_filter(columns: list[dict]) -> list[dict]:
        return [column for column in columns if "relationship" not in column]

    return [
        {**model, "columns": column_filter(model["columns"])} for model in mdl["models"]
    ]

This function filters out existing relationships from the input models, ensuring clean data for new recommendations.

Updated ModelRelationship Class:

class ModelRelationship(BaseModel):
    name: str
    fromModel: str
    fromColumn: str
    type: str
    toModel: str
    toColumn: str
    reason: str

The class now provides more specific and structured information about relationships.

Revised System Prompt:

• Now includes more specific guidelines for relationship recommendations.
• Provides a clearer structure for the output JSON.
• Adds instructions for cases where no relationships should be recommended.

Pipeline Function Updates:

• Added @observe decorators to prompt, generate, and normalized functions.
• Modified prompt function to use the new cleaned_models instead of raw input.

Benefits

1. Improved Accuracy: The new system prompt provides clearer guidelines, potentially leading to more accurate and relevant relationship recommendations.
2. Better Observability: The addition of @observe decorators enhances our ability to monitor and debug the pipeline.
3. Cleaner Input: The cleaned_models function ensures that existing relationships don't interfere with new recommendations.
4. More Detailed Output: The updated ModelRelationship class provides a more comprehensive structure for relationship information.

Testing

This is the testing request body for ensuring the endpoint

{
  "mdl": "{\"models\":[{\"name\":\"Student\",\"columns\":[{\"name\":\"StuID\",\"type\":\"STRING\"},{\"name\":\"Lname\",\"type\":\"STRING\"},{\"name\":\"Fname\",\"type\":\"STRING\"},{\"name\":\"Age\",\"type\":\"INTEGER\"},{\"name\":\"Sex\",\"type\":\"STRING\"},{\"name\":\"Major\",\"type\":\"STRING\"},{\"name\":\"Advisor\",\"type\":\"STRING\"}]},{\"name\":\"Department\",\"columns\":[{\"name\":\"DNO\",\"type\":\"STRING\"},{\"name\":\"DName\",\"type\":\"STRING\"},{\"name\":\"Office\",\"type\":\"STRING\"},{\"name\":\"DPhone\",\"type\":\"STRING\"}]},{\"name\":\"Course\",\"columns\":[{\"name\":\"CrsCode\",\"type\":\"STRING\"},{\"name\":\"DeptNo\",\"type\":\"STRING\"},{\"name\":\"CrsName\",\"type\":\"STRING\"},{\"name\":\"Descr\",\"type\":\"STRING\"},{\"name\":\"Credit\",\"type\":\"INTEGER\"}]},{\"name\":\"Section\",\"columns\":[{\"name\":\"CrsCode\",\"type\":\"STRING\"},{\"name\":\"SecNo\",\"type\":\"STRING\"},{\"name\":\"Semester\",\"type\":\"STRING\"},{\"name\":\"Year\",\"type\":\"INTEGER\"}]},{\"name\":\"Enrolled_in\",\"columns\":[{\"name\":\"StuID\",\"type\":\"STRING\"},{\"name\":\"CrsCode\",\"type\":\"STRING\"},{\"name\":\"SecNo\",\"type\":\"STRING\"},{\"name\":\"Semester\",\"type\":\"STRING\"},{\"name\":\"Year\",\"type\":\"INTEGER\"},{\"name\":\"Grade\",\"type\":\"STRING\"}]},{\"name\":\"Gradeconversion\",\"columns\":[{\"name\":\"lettergrade\",\"type\":\"STRING\"},{\"name\":\"gradepoint\",\"type\":\"FLOAT\"}]},{\"name\":\"Faculty\",\"columns\":[{\"name\":\"FacID\",\"type\":\"STRING\"},{\"name\":\"Lname\",\"type\":\"STRING\"},{\"name\":\"Fname\",\"type\":\"STRING\"},{\"name\":\"Office\",\"type\":\"STRING\"},{\"name\":\"Phone\",\"type\":\"STRING\"},{\"name\":\"Rank\",\"type\":\"STRING\"}]},{\"name\":\"Member_of\",\"columns\":[{\"name\":\"FacID\",\"type\":\"STRING\"},{\"name\":\"DNO\",\"type\":\"STRING\"}]},{\"name\":\"Minor_in\",\"columns\":[{\"name\":\"StuID\",\"type\":\"STRING\"},{\"name\":\"DNO\",\"type\":\"STRING\"}]}],\"relationships\":[{\"name\":\"Student_Enrolled_in\",\"models\":[\"Student\",\"Enrolled_in\"],\"joinType\":\"ONE_TO_MANY\",\"condition\":\"Student.StuID = Enrolled_in.StuID\"},{\"name\":\"Course_Section\",\"models\":[\"Course\",\"Section\"],\"joinType\":\"ONE_TO_MANY\",\"condition\":\"Course.CrsCode = Section.CrsCode\"},{\"name\":\"Section_Enrolled_in\",\"models\":[\"Section\",\"Enrolled_in\"],\"joinType\":\"ONE_TO_MANY\",\"condition\":\"Section.CrsCode = Enrolled_in.CrsCode AND Section.SecNo = Enrolled_in.SecNo AND Section.Semester = Enrolled_in.Semester AND Section.Year = Enrolled_in.Year\"},{\"name\":\"Department_Course\",\"models\":[\"Department\",\"Course\"],\"joinType\":\"ONE_TO_MANY\",\"condition\":\"Department.DNO = Course.DeptNo\"},{\"name\":\"Enrolled_in_Gradeconversion\",\"models\":[\"Enrolled_in\",\"Gradeconversion\"],\"joinType\":\"MANY_TO_ONE\",\"condition\":\"Enrolled_in.Grade = Gradeconversion.lettergrade\"},{\"name\":\"Member_of_Faculty\",\"models\":[\"Member_of\",\"Faculty\"],\"joinType\":\"MANY_TO_ONE\",\"condition\":\"Member_of.FacID = Faculty.FacID\"},{\"name\":\"Member_of_Department\",\"models\":[\"Member_of\",\"Department\"],\"joinType\":\"MANY_TO_ONE\",\"condition\":\"Member_of.DNO = Department.DNO\"},{\"name\":\"Minor_in_Student\",\"models\":[\"Minor_in\",\"Student\"],\"joinType\":\"MANY_TO_ONE\",\"condition\":\"Minor_in.StuID = Student.StuID\"},{\"name\":\"Minor_in_Department\",\"models\":[\"Minor_in\",\"Department\"],\"joinType\":\"MANY_TO_ONE\",\"condition\":\"Minor_in.DNO = Department.DNO\"}],\"metrics\":[],\"cumulativeMetrics\":[],\"enumDefinitions\":[],\"views\":[],\"macros\":[]}"
}

Screenshots

[cyyeh]

LGTM