Update TextGen/router.py

TextGen/router.py

@@ -9,6 +9,7 @@ from fastapi.middleware.cors import CORSMiddleware
 from langchain.chains import LLMChain
 from langchain.prompts import PromptTemplate
 from TextGen.suno import custom_generate_audio, get_audio_information,generate_lyrics
+from TextGen.gemini import generate_story,place_objects,generate_map_markdown
 #from TextGen.diffusion import generate_image
 #from coqui import predict
 from langchain_google_genai import (
@@ -181,10 +182,12 @@ def story(prompt: Invoke):
 def placement(input: Rooms):
     print(input)
     markdown_map=generate_map_markdown(input)
-    print(markdown_map)
-    answer={
-        "key":"value"
-    }
+    story=generate_story(input.possible_entities)
+    print(story)
+    placements=place_objects(input.possible_entities,story,markdown_map)
+    print(placements)
+    answer={"objective":story,
+           "placements":placements}
     return answer
 
 #Dummy function for now
@@ -318,51 +321,4 @@ async def generate_song():
 #    img_byte_arr = img_byte_arr.getvalue()
 #
     # Return the image as a PNG response
-#    return Response(content=img_byte_arr, media_type="image/png")
-
-def generate_map_markdown(data):
-    import numpy as np
-    
-    # Define the room structure with walls and markers
-    def create_room(room_char):
-        return [
-            f"╔═══╗",
-            f"║ {room_char} ║",
-            f"╚═══╝"
-        ]
-    
-    # Extract rooms and rooms of interest
-    rooms = [eval(room) for room in data["rooms"]]
-    rooms_of_interest = [eval(room) for room in data["room_of_interest"]]
-    
-    # Determine grid size
-    min_x = min(room[0] for room in rooms)
-    max_x = max(room[0] for room in rooms)
-    min_y = min(room[1] for room in rooms)
-    max_y = max(room[1] for room in rooms)
-    
-    # Create grid with empty spaces represented by a room-like structure
-    map_height = (max_y - min_y + 1) * 3
-    map_width = (max_x - min_x + 1) * 5
-    grid = np.full((map_height, map_width), " ")
-    
-    # Populate grid with rooms and their characteristics
-    for i, room in enumerate(rooms):
-        x, y = room
-        x_offset = (x - min_x) * 5
-        y_offset = (max_y - y) * 3
-        if room == (0, 0):
-            room_char = "X"
-        elif room in rooms_of_interest:
-            room_char = "P" if i == data["index_exit"] else "?"
-        else:
-            room_char = " "
-        room_structure = create_room(room_char)
-        for j, row in enumerate(room_structure):
-            grid[y_offset + j, x_offset:x_offset + 5] = list(row)
-    
-    # Convert grid to a string format suitable for display
-    markdown_map = "\n".join("".join(row) for row in grid)
-    
-    # Return the map wrapped in triple backticks for proper display in markdown
-    return f"```\n{markdown_map}\n```"
+#    return Response(content=img_byte_arr, media_type="image/png")